Create a TPU instance

This document shows you how to create a TPU v6e (Trillium) virtual machine (VM) instance, connect to it, and run a calculation on it. The TPU instance that you create in this quickstart uses on-demand, the default consumption option.

Before you begin

  1. Install the Google Cloud CLI.

  2. Configure the gcloud CLI to use your federated identity.

    For more information, see Sign in to the gcloud CLI with your federated identity.

  3. To initialize the gcloud CLI, run the following command: 

    gcloud init

  4. Create or select a Cloud de Confiance project.

    Roles required to select or create a project

    • Select a project: Selecting a project doesn't require a specific IAM role—you can select any project that you've been granted a role on.
    • Create a project: To create a project, you need the Project Creator role (roles/resourcemanager.projectCreator), which contains the resourcemanager.projects.create permission. Learn how to grant roles.

    • Create a Cloud de Confiance project:

      gcloud projects create PROJECT_ID

      Replace PROJECT_ID with a name for the Cloud de Confiance project you are creating.

    • Select the Cloud de Confiance project that you created:

      gcloud config set project PROJECT_ID

      Replace PROJECT_ID with your Cloud de Confiance project name.

  5. If you're using an existing project for this guide, verify that you have the permissions required to complete this guide. If you created a new project, then you already have the required permissions.

  6. Verify that billing is enabled for your Cloud de Confiance project.

  7. Enable the Compute Engine API:

    Roles required to enable APIs

    To enable APIs, you need the Service Usage Admin IAM role (roles/serviceusage.serviceUsageAdmin), which contains the serviceusage.services.enable permission. Learn how to grant roles. 

    gcloud services enable compute.googleapis.com

Required roles

To get the permissions that you need to create a TPU instance, connect to it using SSH, and run commands, ask your administrator to grant you the following IAM roles on your project:

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.

Create a TPU v6e (Trillium) instance

Run the following commands in your local terminal or using Cloud Shell.

  1. To create a TPU v6e instance, run the following gcloud compute instances create command:

    gcloud compute instances create quickstart-tpu \
    --machine-type=ct6e-standard-4t \
    --zone=europe-west4-a \
    --image-project=ubuntu-os-accelerator-images \
    --image-family=ubuntu-accel-2204-amd64-tpu-v5e-v5p-v6e \
    --maintenance-policy=TERMINATE

  2. Optionally, verify that the instance is running by using the gcloud compute instances describe command:

    gcloud compute instances describe quickstart-tpu --zone=europe-west4-a

    In the output, check the status field to verify if the instance is running or still pending.

Connect to the TPU instance

You can connect to a TPU instance using the Cloud de Confiance console or the gcloud CLI.

To connect to the VM using SSH, run the following gcloud compute ssh command:

gcloud compute ssh quickstart-tpu --zone=europe-west4-a

Install JAX

Install the JAX framework by running the following command on your TPU instance:

pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html

Run a system check

Verify that JAX can access the TPU and can run basic operations by running the following commands on your TPU instance:

  1. Start the Python 3 interpreter:

    python3

  2. Import the JAX module:

    >>> import jax

  3. Display the number of TPU cores available:

    >>> jax.device_count()

    The output shows the number of TPU cores. The number of cores is dependent on the TPU version you are using. For more information, see TPU versions.

Perform a calculation

  1. Run the following command to add two numbers together:

    >>> jax.numpy.add(1, 1)

    The output should be similar to the following:

    >>> Array(2, dtype=int32, weak_type=True)

  2. Exit the Python interpreter:

    >>> 
exit()

Clean up

To avoid incurring charges to your Cloud de Confiance account for the resources used on this page, delete the Cloud de Confiance project with the resources.

To avoid incurring charges to your Cloud de Confiance by S3NS account for the resources used on this page, delete the TPU instance using the gcloud compute instances delete command:

gcloud compute instances delete quickstart-tpu --zone=europe-west4-a

What's next