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

exportTPU_PROJECT_ID=tpu-project-id
exportSERVICE_PROJECT_ID=service-project-id
exportSERVICE_ACCOUNT_EMAIL=your_service_account@$(SERVICE_PROJECT_ID).iam.gserviceaccount.com
exportZONE=your-zone
exportTPU_NAME=your-tpu-name
exportNETWORK=your-network
exportSUBNET=your-subnet

Set IAM permissions

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

   gcloudprojectsadd-iam-policy-binding${TPU_PROJECT_ID}\
   --member=serviceAccount:${SERVICE_ACCOUNT_EMAIL}\
   --role=roles/cloudtpu.serviceAgent

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

   gcloudprojectsadd-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.

   gcloudiamservice-accountsadd-iam-policy-binding\
   --project$SERVICE_PROJECT_ID$SERVICE_ACCOUNT_EMAIL\
   --roleroles/iam.serviceAccountTokenCreator\
   --memberserviceAccount: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.

gcloudalphacomputetpustpu-vmcreate$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