Create and run a job that uses TPUs

This document explains how to create and run a Batch job that uses virtual machine (VM) instances with Tensor Processing Units (TPUs). TPUs are Google's custom-developed, application-specific integrated circuits (ASICs) designed to accelerate machine learning (ML) and artificial intelligence (AI) workloads. When you create a Batch job, you can optionally specify a TPU machine type to accelerate containerized or script-based AI and ML inference, evaluation, fine-tuning, and training workloads.

To learn more about TPUs, see About TPUs on Google Cloud and TPU machine types in the Compute Engine documentation. Alternatively, you can accelerate data processing, AI, and ML workloads by using GPUs instead.

Before you begin

  1. If you haven't used Batch before, review Get started with Batch and enable Batch by completing the prerequisites for projects and users.
  2. To get the permissions that you need to create a job, ask your administrator to grant you the following IAM roles:

    For more information about granting roles, see Manage access to projects, folders, and organizations.

    You might also be able to get the required permissions through custom roles or other predefined roles.

Restrictions

When a job uses TPUs, all of the following restrictions apply in addition to the Compute Engine limitations for TPUs.

  • Single-host TPU slices only: Batch only supports independent TPU VMs, which are called single-host TPU slices. Multi-host TPU slices, which each consist of multiple interconnected VMs, aren't supported.

    A job can run on one or more single-host TPU slices. For example, if a job sets taskCount to 4 and taskCountPerNode to 1, then the number of VMs the job runs on is 4 or less. Unlike multi-host TPU slices, single-host TPU slices each operate as an independent VM and don't provide inter-chip interconnect (ICI) networking for high-bandwidth, low-latency communication.

    Learn more about single-host and multi-host TPU slices in the Compute Engine documentation.

  • Supported TPU versions and machine types: Batch only supports single-host TPU slices (independent TPU VMs) for the following TPU versions and machine types:

    • TPU7x (Ironwood): tpu7x-standard-4t
    • v6e (Trillium): ct6e-standard-1t, ct6e-standard-4t, and ct6e-standard-8t

    • v5p: ct5p-hightpu-4t

    For more information about each TPU machine type, such as supported disks, see TPU machine types in the Compute Engine documentation.

  • Supported consumption options: Batch doesn't support Flex-start VMs (FLEX_START provisioning model) for jobs that use TPUs. Excluding Flex-start VMs, you can specify any consumption option that is supported by your job's TPU version according to Consumption option availability by TPU versions in the Compute Engine documentation.

  • Single-zone location required: when creating a job that uses TPUs, you must restrict the job to a specific zone by setting the job's allowed locations (allowedLocations[] field) to a zone that contains your TPU version or, if applicable, your reservation.

  • Unsupported features: jobs that use TPUs don't support the following features:

Create a job that uses TPUs

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

  1. Plan the configuration for a job that uses TPUs.
  2. Create a job with the configuration that you identified. For examples of how to create a job that uses TPUs, see Create a job that uses TPUs in this document.

Plan the configuration for a job that uses TPUs

Before you proceed to create a job that uses TPUs, you must do the following:

  1. Choose the TPU version: select the TPU version that is suitable for your workload. For a list of TPU versions by workload type, see Recommended TPU versions by workload types.

    Ensure that you select one of the supported TPU versions.

  2. Choose a TPU consumption option: select a consumption option that best fits your workload, its duration, and your cost needs. Ensure that you select one of the supported consumption options. The use cases and job requirements for each consumption option are summarized in the following table.

    Consumption options and their job requirements

    Reservation-bound (future reservations in calendar mode or future reservations for capacity blocks)

    • Use case: We recommend reservation-bound consumption options for jobs that need a very high level of assurance of TPU resource availability for a specific time period. In Compute Engine, these are described as future reservations in calendar mode (for 1 to 90 days) and future reservations for capacity blocks (for 1 year or longer).

    • Details: Before creating a job, you must obtain or identify a reservation for your chosen TPU version and the RESERVATION_BOUND provisioning model. After a reservation is created, it incurs the costs of the specified VMs at the same price as running the VMs until you delete the reservation. VMs that are consuming a reservation don't incur separate costs, but reservations incur costs regardless of consumption.

      For more information, see Choose a reservation type in the Compute Engine documentation.

    Batch consumes reservation-bound reservations for jobs that do all of the following:

    • Set the provisioningModel field to RESERVATION_BOUND.
    • Specifically target a compatible reservation. Set the reservation field to either the short name (if the reservation is in the same project and zone) or full resource path of the reservation, reservation block, or reservation sub-block.

    On-demand

    • Use case: On-demand is fast and simple to try using, so it's the default consumption option for most jobs. However, we don't recommend on-demand if another consumption option better fits your priorities or if you repeatedly encounter issues with resource obtainability.

    • Details: On-demand is usually the default way to access Compute Engine VMs. On-demand lets you access resources (if available) when requested. The resources aren't deleted until either the job finishes or the run time limit is reached.

    Batch uses on-demand for most jobs by default. If you want to ensure that a job uses on-demand, do all of the following:

    • Set the provisioningModel field to STANDARD.
    • Prevent reservation consumption. To prevent a job from automatically consuming using compatible reservations, set the reservation field to NO_RESERVATION. For more information, see Create and run a job that can't consume reserved VMs.

    Spot VMs

    • Use case: We recommend that you use Spot VMs to reduce costs for fault-tolerant workloads.

    • Details: If a request to create a job that uses Spot VMs is fulfilled, the VMs are charged at Spot VMs pricing, which can change often and might provide the largest discount, up to 91% off of on-demand pricing. However, Spot VMs might not always be available and can be preempted at any time.

      Learn more about Spot VMs and recommendations for jobs with Spot VMs.

    Batch uses Spot VMs for jobs that do all of the following:

    • Set the provisioningModel field to SPOT.
    • (Optional) Prevent reservations consumption. You can't use Spot VMs with reservations. Specifically, the job can optionally set the reservation field to NO_RESERVATION. For more information, see Create and run a job that can't consume reserved VMs.

  3. Select a zone where your chosen TPUs are available: to use a TPU version, a job's allowedLocations[] field must be set to a zone that contains your chosen TPUs as follows:

    • If you chose a consumption option that uses reservations, then you must use the zone where the reservation is located.
    • Otherwise, select a zone that contains your TPU version according to TPU availability.

    In either case, the zone must also be in one of the supported regions for Batch.

  4. Ensure that your project has sufficient TPU quota: if you are creating a job that uses TPUs with on-demand or Spot VMs, you must have sufficient TPU quota available in the region that you want to use. Creating a job that uses TPUs that consumes a TPU reservation doesn't require any TPU quota as quota is used when the reservation was created. For a list of TPU quota names, see TPU quota and for instructions on how to view the quota, see View and manage quotas

  5. Choose a TPU machine type: select a TPU machine type for your chosen TPU version. If you chose a consumption option that uses reservations, then you must use the machine type of your reservation. For a list of machine types available for each TPU version, see TPU machine types.

    Due to the restrictions, you must use the single-host topology for your TPU machine type.

Create a job that uses TPUs

The following sections explain how to create an example job that uses TPUs for each consumption option.

Consumption option Example job instructions
Reservation-bound (future reservations in calendar mode or future reservations for capacity blocks)

Use TPUs with reservation-bound consumption options

On-demand or Spot VMs

Use TPUs with on-demand or Spot VMs

Use TPUs with on-demand or Spot VMs

You can create a job that uses TPUs with on-demand or Spot VMs using the gcloud CLI or Batch API.

gcloud

  1. Create a JSON configuration file that specifies a supported single-host TPU machine type, sets the provisioning model to STANDARD or SPOT, and runs in a location that offers the TPU machine type.

    For example, to create a basic script job, create a JSON file with the following contents:

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

    Replace the following:

    • PROVISIONING_MODEL: the provisioning model for your consumption option: either STANDARD for on-demand or SPOT for Spot VMs.

    • MACHINE_TYPE: a supported single-host TPU machine type—for example, ct5p-hightpu-4t.

    • ALLOWED_LOCATIONS: the allowed locations for the job. For a job that uses TPUs, you must set the allowedLocations[] field to a zone that offers the TPU machine type—for example, zones/us-central1-a offers the TPU machine type ct5p-hightpu-4t.

  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 a supported single-host TPU machine type, sets the provisioning model to STANDARD or SPOT, and runs in a location that offers the TPU machine type.

For example, to create a basic script job, make the following request:

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 TPU task ${BATCH_TASK_INDEX}."
            }
          }
        ]
      },
      "taskCount": 3,
      "parallelism": 1
    }
  ],
  "allocationPolicy": {
    "instances": [
      {
        "policy": {
          "provisioningModel": "PROVISIONING_MODEL",
          "machineType": "MACHINE_TYPE"
        }
      }
    ],
    "location": {
      "allowedLocations": [
        "ALLOWED_LOCATIONS"
      ]
    }
  },
  "logsPolicy": {
    "destination": "CLOUD_LOGGING"
  }
}

Replace the following:

  • PROJECT_ID: the project ID of your project.

  • LOCATION: the location of the job.

  • JOB_NAME: the name of the job.

  • PROVISIONING_MODEL: the provisioning model for your consumption option: either STANDARD for on-demand or SPOT for Spot VMs.

  • MACHINE_TYPE: a supported single-host TPU machine type—for example, ct5p-hightpu-4t.

  • ALLOWED_LOCATIONS: the allowed locations for the job. For a job that uses TPUs, you must set the allowedLocations[] field to a zone that offers the TPU machine type—for example, zones/us-central1-a offers the TPU machine type ct5p-hightpu-4t.

Use TPUs with reservation-bound consumption options

You can create a job that uses TPUs with a consumption option for the reservation-bound (RESERVATION_BOUND) provisioning model—either future reservations for capacity blocks or future reservations in calendar mode—using the gcloud CLI or Batch API.

gcloud

  1. Create a JSON configuration file that specifies a reservation for a single-host TPU machine type, sets the provisioning model to RESERVATION_BOUND, and runs in the zone of the reservation.

    For example, to create a basic script job, create a JSON file with the following contents:

    {
      "taskGroups": [
        {
          "taskSpec": {
            "runnables": [
              {
                "script": {
                  "text": "echo Hello world from TPU task ${BATCH_TASK_INDEX}."
                }
              }
            ]
          },
          "taskCount": 3,
          "parallelism": 1
        }
      ],
      "allocationPolicy": {
        "instances": [
          {
            "policy": {
              "provisioningModel": "RESERVATION_BOUND",
              "reservation": "RESERVATION",
              "machineType": "MACHINE_TYPE"
            }
          }
        ],
        "location": {
          "allowedLocations": [
            "ALLOWED_LOCATIONS"
          ]
        }
      },
      "logsPolicy": {
        "destination": "CLOUD_LOGGING"
      }
    }
    

    Replace the following:

    • RESERVATION: replace with one of the following paths to your reservation:

      • If you want to use a future reservation in calendar mode or any part of the reservation for a future reservation for capacity blocks, then specify the reservation:

        projects/PROJECT_ID/zones/ZONE/reservations/RESERVATION_NAME
        
      • If you want to use a specific block from a future reservation for capacity blocks, then specify the reservation block:

        projects/PROJECT_ID/zones/ZONE/reservations/RESERVATION_NAME/reservationBlocks/BLOCK_NAME
        
      • If you want to use a specific sub-block from a future reservation for capacity blocks, then specify the reservation sub-block:

        projects/PROJECT_ID/zones/ZONE/reservations/RESERVATION_NAME/reservationSubBlocks/SUB_BLOCK_NAME
        
    • MACHINE_TYPE: the TPU machine type of your reservation—for example, ct5p-hightpu-4t.

    • ALLOWED_LOCATIONS: the allowed locations for the job. For a job that uses reserved TPUs, you must set the allowedLocations[] field to the zone of your reservation—for example, zones/us-central1-a offers the TPU machine type ct5p-hightpu-4t.

  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 a reservation for a single-host TPU machine type, sets the provisioning model to RESERVATION_BOUND, and runs in the zone of the reservation.

For example, to create a basic script job, make the following request:

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 TPU task ${BATCH_TASK_INDEX}."
            }
          }
        ]
      },
      "taskCount": 3,
      "parallelism": 1
    }
  ],
  "allocationPolicy": {
    "instances": [
      {
        "policy": {
          "provisioningModel": "RESERVATION_BOUND",
          "reservation": "RESERVATION",
          "machineType": "MACHINE_TYPE"
        }
      }
    ],
    "location": {
      "allowedLocations": [
        "ALLOWED_LOCATIONS"
      ]
    }
  },
  "logsPolicy": {
    "destination": "CLOUD_LOGGING"
  }
}

Replace the following:

  • PROJECT_ID: the project ID of your project.

  • LOCATION: the location of the job.

  • JOB_NAME: the name of the job.

  • RESERVATION: replace with one of the following paths to your reservation:

    • If you want to use a future reservation in calendar mode or any part of the reservation for a future reservation for capacity blocks, then specify the reservation:

      projects/PROJECT_ID/zones/ZONE/reservations/RESERVATION_NAME
      
    • If you want to use a specific block from a future reservation for capacity blocks, then specify the reservation block:

      projects/PROJECT_ID/zones/ZONE/reservations/RESERVATION_NAME/reservationBlocks/BLOCK_NAME
      
    • If you want to use a specific sub-block from a future reservation for capacity blocks, then specify the reservation sub-block:

      projects/PROJECT_ID/zones/ZONE/reservations/RESERVATION_NAME/reservationSubBlocks/SUB_BLOCK_NAME
      
  • MACHINE_TYPE: the TPU machine type of your reservation—for example, ct5p-hightpu-4t.

  • ALLOWED_LOCATIONS: the allowed locations for the job. For a job that uses reserved TPUs, you must set the allowedLocations[] field to the zone of your reservation—for example, zones/us-central1-a offers the TPU machine type ct5p-hightpu-4t.

For more information about creating jobs that consume reservations, see Ensure resource availability using VM reservations.

What's next