Configure networking and access to your Cloud TPU

This page describes how to set up custom network and access configurations for your Cloud TPU, including:

  • Specifying a custom network and subnetwork
  • Enabling SSH access to TPUs
  • Attaching a custom service account to your TPU
  • Enabling custom SSH methods

Prerequisites

Before you run these procedures, you must install the Google Cloud CLI, create a Google Cloud project, and enable the Cloud TPU API. For instructions, see Set up the Cloud TPU environment.

Specify a custom network and subnetwork

You can optionally specify the network and subnetwork to use for the TPU. If the network not specified, the TPU will be in the default network. The subnetwork needs to be in the same region as the zone where the TPU runs.

  1. Create a network that matches one of the following valid formats:

    • https://www.googleapis.com/compute/{version}/projects/{proj-id}/global/networks/{network}
    • compute/{version}/projects/{proj-id}/global/networks/{network}
    • compute/{version}/projects/{proj-##}/global/networks/{network}
    • projects/{proj-id}/global/networks/{network}
    • projects/{proj-##}/global/networks/{network}
    • global/networks/{network}
    • {network}

    For more information, see Create and manage VPC networks.

  2. Create a subnetwork that matches one of the following valid formats:

    • https://www.googleapis.com/compute/{version}/projects/{proj-id}/regions/{region}/subnetworks/{subnetwork}
    • compute/{version}/projects/{proj-id}/regions/{region}/subnetworks/{subnetwork}
    • compute/{version}/projects/{proj-##}/regions/{region}/subnetworks/{subnetwork}
    • projects/{proj-id}/regions/{region}/subnetworks/{subnetwork}
    • projects/{proj-##}/regions/{region}/subnetworks/{subnetwork}
    • regions/{region}/subnetworks/{subnetwork}
    • {subnetwork}

    For more information, see Create and manage VPC networks.

  3. Create the TPU, specifying the custom network and subnetwork:

    gcloud

    To specify the network and subnetwork using the gcloud CLI, add the --network and --subnetwork flags to your create request:

        $ gcloud compute tpus tpu-vm create TPU_NAME \
        --zone=us-central2-b \
        --accelerator-type=v4-8 \
        --version=TPU_SOFTWARE_VERSION
        --network=NETWORK
        --subnetwork=SUBNETWORK

    curl

    To specify the network and subnetwork in a curl call, add the network and subnetwork fields to the request body:

        $ curl -X POST -H "Authorization: Bearer $(gcloud auth print-access-token)" -H "Content-Type: application/json" -d "{accelerator_type: 'v4-8', \
        runtime_version:'tpu-vm-tf-2.17.0-pjrt', \
        network_config: {network: 'NETWORK', subnetwork: 'SUBNETWORK', enable_external_ips: true}, \
        shielded_instance_config: { enable_secure_boot: true }}" \
        https://tpu.googleapis.com/v2/projects/PROJECT_ID/locations/us-central2-b/nodes?node_id=TPU_NAME

  4. Enable custom SSH methods for your network and subnetwork.

Enable SSH access to TPUs

To connect to TPUs using SSH, you need to either enable external IP addresses for the TPUs, or enable Private Google Access for the subnetwork to which the TPU VMs are connected.

Enable external IP addresses

When you create a TPU, external IP addresses are created by default.

If you don't want to create external IP addresses for your TPUs, enable internal IPs:

gcloud

Use the --internal-ips flag when creating a TPU:

$ gcloud compute tpus tpu-vm create TPU_NAME \
    --zone=us-central2-b \
    --accelerator-type=v4-8 \
    --version=TPU_SOFTWARE_VERSION
    --internal-ips

curl

Set the enable_external_ips field to false in the request body:

$ curl -X POST -H "Authorization: Bearer $(gcloud auth print-access-token)" -H "Content-Type: application/json" -d "{accelerator_type: 'v4-8', \
    runtime_version:'tpu-vm-tf-2.17.0-pjrt', \
    network_config: {enable_external_ips: false}, \
    shielded_instance_config: { enable_secure_boot: true }}" \
    https://tpu.googleapis.com/v2/projects/PROJECT_ID/locations/us-central2-b/nodes?node_id=TPU_NAME

Enable Private Google Access

TPUs that don't have external IP addresses can use Private Google Access to access Google APIs and services. For more information about enabling Private Google Access, see Configure Private Google Access.

After you have configured Private Google Access, connect to the VM using SSH.

Attach a custom service account

Each TPU VM has an associated service account it uses to make API requests on your behalf. TPU VMs use this service account to call Cloud TPU APIs and access Cloud Storage and other services. By default, your TPU VM uses the default Compute Engine service account.

The service account must be defined in the same Google Cloud project where you create your TPU VM. Custom service accounts used for TPU VMs must have the TPU Viewer role to call the Cloud TPU API. If the code running in your TPU VM calls other Google Cloud services, it must have the roles necessary to access those services.

For more information about service accounts, see Service accounts.

Use the following commands to specify a custom service account.

gcloud

Use the --service-account flag when creating a TPU:

$ gcloud compute tpus tpu-vm create TPU_NAME \
    --zone=us-central2-b \
    --accelerator-type=TPU_TYPE \
    --version=tpu-vm-tf-2.17.0-pjrt \
    --service-account=SERVICE_ACCOUNT

curl

Set the service_account field in the request body:

$ curl -X POST -H "Authorization: Bearer $(gcloud auth print-access-token)" -H "Content-Type: application/json" -d "{accelerator_type: 'v4-8', \
    runtime_version:'tpu-vm-tf-2.17.0-pjrt', \
    network_config: {enable_external_ips: true}, \
    shielded_instance_config: { enable_secure_boot: true }}" \
    service_account: {email: 'SERVICE_ACCOUNT'} \
    https://tpu.googleapis.com/v2/projects/PROJECT_ID/locations/us-central2-b/nodes?node_id=TPU_NAME

Enable custom SSH methods

The default network allows SSH access to all TPU VMs. If you use a network other than the default or you change the default network settings, you need to explicitly enable SSH access by adding a firewall rule:

$ gcloud compute tpus tpu-vm compute firewall-rules create \
    --network=NETWORK allow-ssh \
    --allow=tcp:22