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@ $( 
SERVICE_PROJECT_ID ) 
.iam.gserviceaccount.com 
 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_EMAIL 
    } 
     
    \ 
   --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