Use a cross-project service account

This document outlines how to use a cross-project service account when creating a TPU VM. A cross-project service account is a service account that is located in a different project than the TPU VM.

Note, the following instructions don't take Shared VPC or VPC Service Controls into consideration. Additional configurations or permissions may be required to support their use. For more information, see Create and modify Shared VPC networks and Overview of VPC Service Controls.

Before you begin

• Create two Google Cloud projects. One to contain the TPU VM (tpu_project) and one to contain the service account (service_project). For more information, see Creating and managing projects.
• Create a service account in the service_project. For more information, see Creating and managing service accounts.
• Create a network and enable Private Google Access. For more information, see Creating and managing networks and Configure Private Google Access for an internal IP address.
• Disable the following organization policies. For more information, see Organization policy constraints.
  • constraints/iam.disableCrossProjectServiceAccountUsage in both projects
  • constraints/compute.vmExternalIpAccess (for accessing the VM with SSH, optional)

Define environment variables

export TPU_PROJECT_ID=tpu-project-id
export SERVICE_PROJECT_ID=service-project-id
export SERVICE_ACCOUNT_EMAIL=your-service-account-email
export ZONE=your-zone
export TPU_NAME=your-tpu-name
export NETWORK=your-network
export SUBNET=your-subnet

Set IAM permissions

1. Allow your service account to use the TPU Service Agent role in the tpu_project.

   gcloud projects add-iam-policy-binding ${TPU_PROJECT_ID} \
   --member=serviceAccount:${SERVICE_ACCOUNT_ID} \
   --role=roles/cloudtpu.serviceAgent

2. Add the serviceAccountUser role to the TPU Service Agent.

   gcloud projects add-iam-policy-binding $SERVICE_PROJECT_ID \
   --member=serviceAccount:service-$TPU_PROJECT_ID@gcp-sa-tpu.iam.gserviceaccount.com \
   --role=roles/iam.serviceAccountUser

3. If your workload is running in Docker containers, allow the Compute Engine Service Agent to access the metadata server.

   gcloud iam service-accounts add-iam-policy-binding  \
   --project $SERVICE_PROJECT_ID $SERVICE_ACCOUNT_EMAIL \
   --role roles/iam.serviceAccountTokenCreator \
   --member serviceAccount:service-$TPU_PROJECT_ID@compute-system.iam.gserviceaccount.com

Create a TPU VM in the tpu_project

Create a TPU VM in the tpu_project using the service account in the service_project.

gcloud alpha compute tpus tpu-vm create $TPU_NAME \
--description=$TPU_NAME \
--accelerator-type=v5litepod-8 \
--version=tpu-vm-tf-2.17.0-pod \
--network=projects/$PROJECT_A/global/networks/$NETWORK \
--subnetwork=$SUBNET \
--internal-ips \
--service-account=$SERVICE_ACCOUNT_EMAIL \
--project=$TPU_PROJECT_ID \
--zone=$ZONE