Run a small batch workload with TPUs and Flex-start VMs

gcloud

To create a multi-host TPU slice node pool with Ironwood (TPU7x), you must first create a workload policy.

1. Create a workload policy:

   gcloud compute resource-policies create workload-policy WORKLOAD_POLICY_NAME \
       --type=HIGH_THROUGHPUT \
       --accelerator-topology=TPU_TOPOLOGY \
       --project=PROJECT_ID \
       --region=REGION
   

   Replace the following:

   • WORKLOAD_POLICY_NAME: a name for your workload policy.
   • TPU_TOPOLOGY: the TPU Ironwood (TPU7x) topology. For example, 2x2x2. To see all supported Ironwood (TPU7x) topologies, see the topology section.
   • PROJECT_ID: your Google Cloud project ID.
   • REGION: the region for the workload policy. A workload policy is a regional resource and can be re-used across node pools that share the same topology.

2. Create the node pool with the workload policy:

   gcloud container node-pools create NODE_POOL_NAME \
       --cluster=CLUSTER_NAME \
       --location=us-central1 \
       --node-locations=us-central1-c \
       --machine-type=tpu7x-standard-4t \
       --reservation-affinity=none \
       --enable-autoscaling \
       --num-nodes=0 --min-nodes=0 --max-nodes=MAX_NODES \
       --flex-start \
       --placement-policy=WORKLOAD_POLICY
   

   Replace the following:

   • NODE_POOL_NAME: the name of the new node pool.
   • WORKLOAD_POLICY: the name of the workload policy that you created.
   • MAX_NODES: The maximum size of the node pool. The --max-nodes flag is required if --enable-autoscaling is supplied and must be equal to the product of the values defined in TPU_TOPOLOGY ({A}x{B}x{C}) divided by the number of chips in each VM. For example, if TPU_TOPOLOGY is 2x2x2, the product is 8. Since each VM in tpu7x-standard-4t has 4 chips, the number of nodes is 2.

   This command creates a node pool named NODE_POOL_NAME with the following characteristics:

   • --machine-type=tpu7x-standard-4t specifies the Ironwood (TPU7x) machine type.
   • --flex-start enables flex-start.

Terraform

1. Ensure that you use the version 4.84.0 or later of the google provider.

2. Create a workload policy:

   resource "google_compute_resource_policy" {
     name   = "WORKLOAD_POLICY_NAME"
     region = CLUSTER_LOCATION
     workload_policy {
       type = "HIGH_THROUGHPUT"
       accelerator_topology = "TPU_TOPOLOGY"
     }
   }
   

   Replace the following:

   • WORKLOAD_POLICY_NAME: a name for your workload policy.
   • CLUSTER_LOCATION: Compute location for the cluster. We recommend having a regional cluster for higher reliability of the Kubernetes control plane. You can also use a zonal cluster. For more information, see Select a TPU version and topology.
   • TPU_TOPOLOGY: the TPU Ironwood (TPU7x) topology. For example, 2x2x2. To see all supported Ironwood (TPU7x) topologies, see Plan TPUs.

   For more information about the google_compute_resource_policy reference, see Terraform Provider.

3. In your Terraform configuration, add the following block:

   resource "google_container_node_pool" "NODE_POOL_RESOURCE_NAME" {
     provider           = google
     project            = PROJECT_ID
     cluster            = CLUSTER_NAME
     name               = POOL_NAME
     location           = CLUSTER_LOCATION
     node_locations     = [NODE_ZONES]
     initial_node_count = NUM_NODES
   
     autoscaling {
       max_node_count = MAX_NODES
       location_policy      = "ANY"
     }
     node_config {
       machine_type = MACHINE_TYPE
       reservation_affinity {
         consume_reservation_type = "SPECIFIC_RESERVATION"
         key = "compute.googleapis.com/reservation-name"
         values = [RESERVATION_LABEL_VALUES]
       }
       flex_start = true
     }
   
     placement_policy {
       policy_name = WORKLOAD_POLICY_NAME
     }
   }
   

   Replace the following:

   • NODE_POOL_RESOURCE_NAME: the name of the node pool resource in the Terraform template.
   • PROJECT_ID: your project ID.
   • CLUSTER_NAME: the name of the existing cluster to add the node pool to.
   • POOL_NAME: the name of the node pool to create.
   • NODE_ZONES: the comma-separated list of one or more zones where GKE creates the node pool.
   • NUM_NODES: the number of nodes in the node pool. It must be zero or the product of the number of the TPU chips divided by four, because in multi-host TPU slices each TPU slice node has four chips. For example, if TPU_TOPOLOGY is 4x8, then there are 32 chips, which means NUM_NODES must be 8. To learn more about TPU topologies, use the table in Choose the TPU version.
   • TPU_TOPOLOGY: this indicates the selected physical topology for the TPU slice. The format of the topology depends on the TPU version you are using. To learn more about TPU topologies, use the table in Choose a topology.

   Optionally, you can also use the following variables:

   • RESERVATION_NAME: if you use a TPU reservation, provide a list of reservation-resource labels to use when creating the node pool. To learn more about how to populate theRESERVATION_LABEL_VALUES in the reservation_affinity field, see Terraform Provider.
   • autoscaling: create a node pool with autoscaling enabled. When GKE scales a multi-host TPU slice node pool, it atomically scales up the node pool from zero to the maximum size.
     • MAX_NODES: the maximum size of the node pool. The value must be equal to the product of the values defined in TPU_TOPOLOGY ({A}x{B}x{C}) divided by the number of chips in each VM. For example, if TPU_TOPOLOGY is 2x2x2, the product is 8. Since each VM in tpu7x-standard-4t has 4 chips, the number of nodes is 2.
   • spot: the node pool that will use Spot VMs for the TPU slice nodes. This setting cannot be changed after the node pool is created. For more information, see Spot VMs.
   • flex_start: the node pool that will use flex-start consumption option. This setting can't be set to true if spot is enabled.

gcloud

  gcloud container node-pools create NODE_POOL_NAME \
      --cluster=CLUSTER_NAME \
      --location=CONTROL_PLANE_LOCATION \
      --node-locations=NODE_ZONES \
      --machine-type=MACHINE_TYPE \
      --tpu-topology=TPU_TOPOLOGY \
      --reservation-affinity=none \
      --enable-autoscaling \
      --num-nodes=0 --min-nodes=0 --max-nodes=MAX_NODES \
      --flex-start

Replace the following:

• NODE_POOL_NAME: the name of the new node pool.
• CLUSTER_NAME: the name of the cluster.
• CONTROL_PLANE_LOCATION: the name of the zone based on the TPU version you want to use. To identify an available location, see TPU availability in GKE.
• NODE_ZONES: the comma-separated list of one or more zones where GKE creates the node pool.
• MACHINE_TYPE: the type of machine to use for nodes. For more information about TPU compatible machine types, use the table in Choose the TPU version.
• TPU_TOPOLOGY: the TPU topology. For example, 2x2x2. To see all supported TPU topologies, see the topology section.
• MAX_NODES: the maximum size of the node pool. The --max-nodes flag is required if --enable-autoscaling is supplied and must be equal to the product of the values defined in TPU_TOPOLOGY ({A}x{B}x{C}) divided by the number of chips in each VM.

  This command creates a node pool named NODE_POOL_NAME with flex-start enabled.

Terraform

1. Ensure that you use the version 4.84.0 or later of the google provider.

2. Add the following block to your Terraform configuration:

   resource "google_container_node_pool" "NODE_POOL_RESOURCE_NAME" {
     provider           = google
     project            = PROJECT_ID
     cluster            = CLUSTER_NAME
     name               = POOL_NAME
     location           = CLUSTER_LOCATION
     node_locations     = [NODE_ZONES]
     initial_node_count = NUM_NODES
   
     autoscaling {
       max_node_count = MAX_NODES
       location_policy      = "ANY"
     }
     node_config {
       machine_type = MACHINE_TYPE
       reservation_affinity {
         consume_reservation_type = "SPECIFIC_RESERVATION"
         key = "compute.googleapis.com/reservation-name"
         values = [RESERVATION_LABEL_VALUES]
       }
       flex_start = true
     }
   
     placement_policy {
       type = "COMPACT"
       tpu_topology = TPU_TOPOLOGY
     }
   }
   

   Replace the following:

   • NODE_POOL_RESOURCE_NAME: the name of the node pool resource in the Terraform template.
   • PROJECT_ID: your project ID.
   • CLUSTER_NAME: the name of the existing cluster to add the node pool to.
   • POOL_NAME: the name of the node pool to create.
   • CLUSTER_LOCATION: compute location for the cluster. We recommend having a regional cluster for higher reliability of the Kubernetes control plane. You can also use a zonal cluster. To learn more, see Select a TPU version and topology.
   • NODE_ZONES: the comma-separated list of one or more zones where GKE creates the node pool.
   • NUM_NODES: the number of nodes in the node pool. It must be zero or the product of the number of the TPU chips divided by four, because in multi-host TPU slices each TPU slice node has 4 chips. For example, if TPU_TOPOLOGY is 4x8, then there are 32 chips which means NUM_NODES must be 8. To learn more about TPU topologies, use the table in Choose the TPU version.
   • TPU_TOPOLOGY: this indicates the physical topology for the TPU slice. The format of the topology depends on the TPU version you are using. To learn more about TPU topologies, use the table in Choose a topology.

   Optionally, you can also use the following variables:

   • RESERVATION_NAME: if you use TPU reservation, this is the list of labels of the reservation resources to use when creating the node pool. To learn more about how to populate theRESERVATION_LABEL_VALUES in the reservation_affinity field, see Terraform Provider.
   • autoscaling: Create a node pool with autoscaling enabled. When GKE scales a multi-host TPU slice node pool, it atomically scales up the node pool from zero to the maximum size.
     • MAX_NODES: it is the maximum size of the node pool. It must be equal to the product of the values defined in TPU_TOPOLOGY ({A}x{B}x{C}) divided by the number of chips in each VM).
   • spot: lets the node pool to use Spot VMs for the TPU slice nodes. This cannot be changed after node pool creation. For more information, see Spot VMs.
   • flex_start: Sets the node pool to use flex-start consumption option. Can't be set to true if spot is enabled.

Console

To create a node pool with TPUs:

1. Go to the Google Kubernetes Engine page in the Google Cloud console.

   Go to Google Kubernetes Engine

2. In the cluster list, click the name of the cluster you want to modify.

3. Click add_box Add node pool.

4. In the Node pool details section, check the Specify node locations box.

5. Select the name of the zone based on the TPU version you want to use. To identify an available location, see TPU availability in GKE.

6. From the navigation pane, click Nodes.

7. In the Machine Configuration section, select TPUs.

8. In the Series drop-down menu, select one of the following:

   • CT3: TPU v3, single-host device
   • CT3P: TPU v3, multi-host pod slice
   • CT4P: TPU v4
   • CT5LP: TPU v5e
   • CT5P: TPU v5p
   • CT6E: TPU Trillium (v6e)

9. In the Machine type drop-down menu, select the name of the machine to use for nodes. Use the Choose the TPU version table to learn how to define the machine type and TPU topology that create a multi-host TPU slice node pool.

10. In the TPU Topology drop-down menu, select the physical topology for the TPU slice.

11. In the Changes needed dialog, click Make changes.

12. Ensure that Boot disk type is either Standard persistent disk or SSD persistent disk.

13. Optionally, select the Enable nodes on spot VMs checkbox to use Spot VMs for the nodes in the node pool.

14. Click Create.