Crea y ejecuta un trabajo que use GPU

En este documento, se explica cómo crear y ejecutar un trabajo que use una unidad de procesamiento de gráficos (GPU). Para obtener más información sobre las funciones y las restricciones de las GPUs, consulta Acerca de las GPUs en la documentación de Compute Engine.

Cuando creas un trabajo por lotes, tienes la opción de usar las GPU para acelerar cargas de trabajo específicas. Entre los casos de uso comunes para los trabajos que usan GPU, se incluyen los siguientes: procesamiento de datos intensivo y cargas de trabajo de inteligencia artificial (IA), como el aprendizaje automático (AA).

Antes de comenzar

  1. Si nunca antes usaste Batch, consulta Cómo comenzar a usar Batch y habilita Batch completando los requisitos previos para proyectos y usuarios.
  2. Para obtener los permisos que necesitas para crear un trabajo, pídele a tu administrador que te otorgue los siguientes roles de IAM:

    Para obtener más información sobre cómo otorgar roles, consulta Administra el acceso a proyectos, carpetas y organizaciones.

    También puedes obtener los permisos necesarios mediante roles personalizados o cualquier otro rol predefinido.

Crea un trabajo que use GPUs

Para crear un trabajo que use GPUs, haz lo siguiente:

  1. Planifica los requisitos de un trabajo que use GPUs.
  2. Crea un trabajo con los requisitos y métodos que identificaste. Para ver ejemplos de cómo crear un trabajo con los métodos recomendados, consulta Crea un trabajo de ejemplo que use GPU en esta .

Planifica los requisitos de un trabajo que use GPUs

Antes de crear un trabajo que use GPUs, planifica los requisitos del trabajo como se explica en las siguientes secciones:

  1. Selecciona el tipo de máquina de la GPU
  2. Instala los controladores de GPU
  3. Define los recursos de VM compatibles

Paso 1: Selecciona el tipo de máquina de la GPU

Los tipos de máquinas de GPU disponibles (las combinaciones válidas de tipos de GPU, cantidades de GPU y tipo de máquina (CPU virtuales y memoria)) y sus casos de uso se enumeran en el La página Tipos de máquinas de GPU en Compute Engine en la documentación de Google Cloud.

Los campos obligatorios para que un trabajo especifique un tipo de máquina de GPU varían según la categorías en la siguiente tabla:

Tipos de máquinas de GPU y requisitos de trabajo

GPUs para VMs optimizadas para aceleradores: Las VMs con un tipo de máquina de la familia de máquinas optimizadas para aceleradores tienen un tipo y una cantidad específicos de estas GPUs conectadas automáticamente.

Para usar GPUs en VMs optimizadas para aceleradores, te recomendamos que especifiques el tipo de máquina. Cada tipo de máquina optimizada para aceleradores solo admite un tipo y una cantidad específicos de GPUs, por lo que es funcionalmente equivalente, ya sea que especifiques esos valores además del tipo de máquina optimizada para aceleradores.

En concreto, Batch también admite especificar solo tipo y cantidad de GPU para las VMs con optimización de acelerador, pero el resultado Las opciones de CPU virtual y memoria suelen ser muy limitadas. Por eso, recomendamos verificas que las opciones disponibles de CPU virtual y memoria sean compatibles con los requisitos de la tarea del trabajo.

GPU para VMs N1: Estas GPU requieren que especifiques el tipo y la cantidad para adjuntar a cada VM y debe conectarse a las VMs con una escribir desde Serie de máquinas N1.

Si quieres usar GPU para las VMs N1, te recomendamos que especifiques al menos la el tipo de GPU y la cantidad de GPU. Asegúrate de que la combinación de de salida coincida con uno de los valores válidos Opciones de GPU para los tipos de máquina N1. Las opciones de CPU virtual y memoria para las VMs N1 que usan cualquier tipo específico GPU es bastante flexible, así que, si lo prefieres, puedes dejar que Batch seleccione una que cumple con los requisitos de la tarea del trabajo.

Paso 2: Instala los controladores de GPU

Para instalar los controladores de GPU requeridos, selecciona uno de los siguientes métodos:

Paso 3: Define los recursos de VM compatibles

Para obtener información sobre los requisitos y las opciones para definir los recursos de VM de un trabajo, consulta Recursos de trabajo.

En resumen, debes hacer todo lo siguiente cuando definas los recursos de VM para un trabajo que usa GPU:

  • Asegúrate de que el tipo de máquina de GPU esté disponible en la ubicación de las VMs de tu trabajo.

    Para obtener información sobre dónde están disponibles los tipos de máquinas con GPU, consulta Disponibilidad de GPU por regiones y zonas en la documentación de Compute Engine.

  • Si especificas el tipo de máquina del trabajo, asegúrate de que tenga suficientes CPUs y memoria para los requisitos de la tarea del trabajo. Especificar el tipo de máquina del trabajo es obligatorio cuando se usan GPUs para VMs optimizadas para aceleradores y opcional cuando se usan GPUs para VMs N1.

  • Asegúrate de definir los recursos de VM para un trabajo con un método válido:

    • Define los recursos de la VM directamente con el campo instances[].policy (recomendado si es posible). Este método se muestra en los ejemplos.
    • Define los recursos de VM a través de una plantilla con el campo instances[].instanceTemplate. Este método es necesario para instalar controladores de GPU de forma manual a través de un imagen. Para obtener más información, consulta Define los recursos de trabajo con una plantilla de instancia de VM.

Crea un trabajo de ejemplo que use GPUs

En las siguientes secciones, se explica cómo crear un trabajo de ejemplo para diferentes tipos de máquinas de GPU. Los trabajos de ejemplo instalan controladores de GPU automáticamente y definir directamente los recursos de VM.

Usa GPU para las VMs con optimización de acelerador

Puedes crear un trabajo que use GPU para VMs con optimización de acelerador mediante el gcloud CLI, API de Batch, Java o Python.

gcloud

  1. Crea un archivo JSON que instale controladores de GPU, defina el campo machineType con el tipo de máquina de la familia de máquinas optimizadas para aceleradores y use una ubicación que tenga el tipo de GPU especificado.

    Por ejemplo, para crear un trabajo de secuencia de comandos básico que use GPU para con las VMs optimizadas para aceleradores, crea un archivo JSON con los siguientes contenidos:

    {
        "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"
                ]
            }
        }
    }
    

    Reemplaza lo siguiente:

    • INSTALL_GPU_DRIVERS: Cuando se establece en true, Batch recupera los controladores necesarios para la Tipo de GPU que especificas en el campo policy desde una ubicación de terceros y Batch las instala por ti. Si estableces este campo en false (configuración predeterminada), debes instalar los controladores de GPU de forma manual para usar cualquier GPU en esta tarea.

    • MACHINE_TYPE: Es un tipo de máquina del familia de máquinas con optimización de acelerador.

    • ALLOWED_LOCATIONS: El Campo allowedLocations[] define una región y, opcionalmente, una o más zonas, donde la VM que tus instancias de trabajo pueden ejecutarse ejemplo, regions/us-central1, zones/us-central1-a permite que la zona us-central1-a Asegúrate de especificar las ubicaciones que ofrecen el tipo de máquina con GPU que deseas para esta tarea. De lo contrario, si omites este campo, y asegurarte de que la ubicación del trabajo ofrezca el tipo de máquina de la GPU.

  2. Para crear y ejecutar el trabajo, usa el Comando gcloud batch jobs submit:

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

    Reemplaza lo siguiente:

    • JOB_NAME: Es el nombre del trabajo.

    • LOCATION: Es la ubicación del trabajo.

    • JSON_CONFIGURATION_FILE: Es la ruta de acceso a un archivo JSON con los detalles de configuración de la tarea.

API

Realiza una solicitud POST al método jobs.create que instale los controladores de GPU, defina el campo machineType con el tipo de máquina de la familia de máquinas optimizadas para aceleradores y use una ubicación que tenga el tipo de GPU especificado.

Por ejemplo, para crear un trabajo de secuencia de comandos básico que use GPU para con las VMs optimizadas para aceleradores, realiza la siguiente solicitud:

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"
            ]
        }
    }
}

Reemplaza lo siguiente:

  • PROJECT_ID: el ID del proyecto de tu proyecto.

  • LOCATION: Es la ubicación del trabajo.

  • JOB_NAME: Es el nombre del trabajo.

  • INSTALL_GPU_DRIVERS: Cuando se establece en true, Batch recupera los controladores necesarios para el tipo de GPU que especificas en el campo policy desde una ubicación de terceros y los instala en tu nombre. Si configuras este campo como false (predeterminado), necesitas para instalar controladores de GPU de forma manual y usar cualquier GPU para este trabajo.

  • MACHINE_TYPE: Es un tipo de máquina del familia de máquinas con optimización de acelerador.

  • ALLOWED_LOCATIONS: El campo allowedLocations[] define una región y, de manera opcional, una o más zonas, en las que se pueden ejecutar las instancias de VM de tu trabajo. Por ejemplo, regions/us-central1, zones/us-central1-a permite la zona us-central1-a. Asegúrate de especificar las ubicaciones que ofrecen el tipo de máquina con GPU que deseas para esta tarea. De lo contrario, si omites este campo, y asegurarte de que la ubicación del trabajo ofrezca el tipo de máquina de la 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)

Usa GPUs para VMs N1

Puedes crear un trabajo que use GPUs para VMs N1 con gcloud CLI, la API de Batch, Java, Node.js o Python.

gcloud

  1. Crea un archivo JSON que instale controladores de GPU, defina los subcampos type y count del campo accelerators[] y use una ubicación que tenga el tipo de GPU especificado.

    Por ejemplo, para crear un trabajo de secuencia de comandos básico que use GPU para VMs N1 y permite que Batch seleccione el tipo de máquina N1 exacto, crea un archivo JSON con el siguiente contenido:

    {
        "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"
                ]
            }
        }
    }
    

    Reemplaza lo siguiente:

    • INSTALL_GPU_DRIVERS: Cuando se establece en true, Batch recupera los controladores necesarios para la Tipo de GPU que especificas en el campo policy desde una ubicación de terceros y Batch las instala por ti. Si estableces este campo en false (configuración predeterminada), debes instalar los controladores de GPU de forma manual para usar cualquier GPU en esta tarea.

    • GPU_TYPE: La GPU el tipo de letra. Puedes ver una lista de los tipos de GPU disponibles con el comando gcloud compute accelerator-types list. Usa este campo solo para las GPU de las VMs N1.

    • GPU_COUNT: Es la cantidad de GPU de la el tipo especificado. Para obtener más información sobre las opciones válidas, consulta los tipos de máquinas de GPU para la serie de máquinas N1. Usa este campo solo para las GPUs de las VMs N1.

    • ALLOWED_LOCATIONS: El Campo allowedLocations[] define una región y, opcionalmente, una o más zonas, donde la VM que tus instancias de trabajo pueden ejecutarse ejemplo, regions/us-central1, zones/us-central1-a permite que la zona us-central1-a Asegúrate de especificar las ubicaciones que ofrecen el tipo de máquina con GPU que deseas para esta tarea. De lo contrario, si omites este campo, asegúrate de que la ubicación del trabajo ofrezca el tipo de máquina de GPU.

  2. Para crear y ejecutar la tarea, usa el comando gcloud batch jobs submit:

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

    Reemplaza lo siguiente:

    • JOB_NAME: Es el nombre del trabajo.

    • LOCATION: Es la ubicación del trabajo.

    • JSON_CONFIGURATION_FILE: Es la ruta de acceso para un Archivo JSON con los detalles de configuración del trabajo.

API

Realiza una solicitud POST al Método jobs.create que instala controladores de GPU, define la Subcampos type y count del campo accelerators[] y usa una ubicación que tiene la ubicación el tipo de GPU.

Por ejemplo, para crear un trabajo de secuencia de comandos básico que use GPUs para VMs N1 y permita que Batch seleccione el tipo de máquina N1 exacto, realiza la siguiente solicitud:

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"
            ]
        }
    }
}

Reemplaza lo siguiente:

  • PROJECT_ID: Es el ID del proyecto de tu proyecto.

  • LOCATION: Es la ubicación del trabajo.

  • JOB_NAME: Es el nombre del trabajo.

  • INSTALL_GPU_DRIVERS: Cuando se establece en true, Batch recupera los controladores necesarios para la Tipo de GPU que especificas en el campo policy desde una ubicación de terceros y Batch las instala por ti. Si configuras este campo como false (predeterminado), necesitas para instalar controladores de GPU de forma manual y usar cualquier GPU para este trabajo.

  • GPU_TYPE: La GPU el tipo de letra. Puedes ver una lista de los tipos de GPU disponibles con el comando Comando gcloud compute accelerator-types list Usa este campo solo para las GPU de las VMs N1.

  • GPU_COUNT: Es la cantidad de GPU de la el tipo especificado. Para obtener más información sobre las opciones válidas, consulta Tipos de máquinas de GPU para la serie de máquinas N1. Usa este campo solo para las GPUs de las VMs N1.

  • ALLOWED_LOCATIONS: El campo allowedLocations[] define una región y, de manera opcional, una o más zonas, en las que se pueden ejecutar las instancias de VM de tu trabajo. Por ejemplo, regions/us-central1, zones/us-central1-a permite la zona us-central1-a. Asegúrate de especificar las ubicaciones que ofrecen el tipo de máquina con GPU que deseas para esta tarea. De lo contrario, si omites este campo, asegúrate de que la ubicación del trabajo ofrezca el tipo de máquina de 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 crear un trabajo con GPUs usando Node.js, selecciona una de las siguientes opciones según el tipo de máquina de tu modelo de GPU:

Crea un trabajo que use GPUs con VMs optimizadas para aceleradores

Para usar GPU con VMs optimizadas para aceleradores, solo especifica el tipo de máquina que quieres para las VMs del trabajo:

// 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();

Crear un trabajo que use GPU con VMs N1

Para usar GPU con VMs N1, debes especificar la cantidad y el tipo de GPU que deseas para cada una de las VMs del trabajo:

// 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)

¿Qué sigue?