Console
Start free

TPU collection scheduling for inference workloads

Trillium (v6e) includes a feature called "collection scheduling" that lets you group a set of TPU slices, single or multi-host, intended to serve replicas of the same model. This feature is available for both Cloud TPU and Google Kubernetes Engine (GKE) configurations.

This document is about using collection scheduling with the Cloud TPU API. See the GKE documentation for more information about using collection scheduling with GKE.

By creating a collection for your inference workload, Google Cloud limits and streamlines interruptions to the operations of inference workloads. This is useful for inference workloads where high availability is a concern. Google Cloud ensures high availability for the collection to manage incoming traffic. A portion of slices within a collection is always available to handle incoming traffic.

Each TPU slice in a collection will have the same accelerator type and topology.

Create a collection from the Cloud TPU API

When you request a queued resource using the Cloud TPU API, you use the --workload-type=AVAILABILITY-OPTIMIZED flag to create a collection. This flag indicates to the Cloud TPU infrastructure that it is meant to be used for availability-focused workloads.

The following command provisions a collection using the Cloud TPU API:

gcloud alpha compute tpus queued-resources create QUEUED_RESOURCE_ID \
   --project=PROJECT_ID \
   --zone=ZONE \
   --accelerator-type=ACCELERATOR_TYPE \
   --runtime-version=RUNTIME_VERSION \
   --node-count=NODE_COUNT \
   --node-prefix=NODE_PREFIX \
   --workload-type=AVAILABILITY-OPTIMIZED

The --node-count flag specifies the number of slices you want in your queued resource. This creates a collection of TPU slices.

Optional: The --node-prefix flag specifies a prefix for the slice names.

Note: The supported accelerator types are described in v6e supported configurations.