Create and run a job that uses GPUs

This document explains how to create and run a job that uses a graphics processing unit (GPU).

When you create a Batch job, you can optionally add one or more GPUs to the VMs running it. Common use cases for jobs that use GPUs include intensive data processing and machine learning (ML) workloads.

Before you begin

Create a job that uses GPUs

To create a job that uses GPUs, do the following:

  1. Review the Requirements for a job to use GPUs section to determine the methods you can use to create your job.
  2. Create a job with the methods you selected. For examples of how to create a job using the recommended methods, see the Create an example job that uses GPUs section.

Requirements for a job to use GPUs

To use GPUs, a job must do all of the following:

After you've determined how to meet these requirements for your job, you also need to define job's GPUs and location. A job's VMs can each use one or more GPUs of the type that you specify. The allowed locations for the job's VMs (or if undefined, the location of the job) must have the specified type of GPUs. For more information about defining the GPU type, GPU number, and a valid location for a job, see the examples.

Install GPU drivers

To install the required GPU drivers, select one of the following methods:

Define compatible VM resources

If your job defines any of the VM resources (any of the instances[] subfields) other than GPUs, you must define those VM resources in a compatible way.

To define the resources for a job's VMs, including any GPUs, you can only use of the following methods:

  • Define resources directly (recommended): As shown in the examples, to define the resources for a job's VMs directly, use the policy field.
  • Define resources in a template: Define the resource's for a job's VMs by specifying a Compute Engine instance template.

Additionally, all of the resources that you define must be compatible with the type and number of the GPUs for the job. For more information about the VM resources that you can use with GPUs, see GPU platforms in the Compute Engine documentation.

Create an example job that uses GPUs

You can create a job that uses GPUs using the gcloud CLI or Batch API.

gcloud

  1. Create a JSON file that specifies the job's configuration details, the type and count subfields of the accelerators[] field, and a location that has those types of GPUs.

    For example, to create a basic script job that uses GPUs, automatically installs the required GPU drivers, and specifies the allowed locations for the job's VMs, create a JSON file with the following contents:

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

    Replace the following:

    • INSTALL_GPU_DRIVERS: 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.

    • GPU_TYPE: the GPU type. You can view a list of the available GPU types by using the gcloud compute accelerator-types list command.

    • GPU_COUNT: the number of GPUs of the specified type.

    • ALLOWED_LOCATIONS: Optional. The locations where the VM instances for your job are allowed to run—for example, regions/us-central1, zones/us-central1-a allows the zone us-central1-a. If you specify an allowed location, you must select the region and, optionally, one or more zones. The locations that you choose must have the GPU types you want for this job. Otherwise, if you omit this field, the job's location must have the GPU types. For more information, see the allowedLocations[] field.

  2. To create and run the job, use the gcloud batch jobs submit command:

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

    Replace the following:

    • JOB_NAME: the name of the job.

    • LOCATION: the location of the job.

    • JSON_CONFIGURATION_FILE: the path for a JSON file with the job's configuration details.

API

Make a POST request to the jobs.create method that specifies the job's configuration details, the type and count subfields of the accelerators[] field, and a location that has those types of GPUs.

For example, to create a basic script job that uses GPUs, automatically installs the required GPU drivers, and specifies the allowed locations for the job's VMs, create a JSON file with the following contents:

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

Replace the following:

  • PROJECT_ID: the project ID of your project.

  • LOCATION: the location of the job.

  • JOB_NAME: the name of the job.

  • INSTALL_GPU_DRIVERS: 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.

  • GPU_TYPE: the GPU type. You can view a list of the available GPU types by using the gcloud compute accelerator-types list command.

  • GPU_COUNT: the number of GPUs of the specified type.

  • ALLOWED_LOCATIONS: Optional. The locations where the VM instances for your job are allowed to run—for example, regions/us-central1, zones/us-central1-a allows the zone us-central1-a. If you specify an allowed location, you must select the region and, optionally, one or more zones. The locations that you choose must have the GPU types you want for this job. Otherwise, if you omit this field, the job's location must have the GPU types. For more information, see the allowedLocations[] field.

What's next