Use Identity Aware Proxy to tunnel to a TPU

user_795053 0
Viewed 241

Can I use google cloud's identity aware proxy to connect to the gRPC endpoint on a TPU worker? By "TPU worker" I mean that I am creating a TPU with no associated compute instance (using gcloud compute tpus create) and I wish to connect to the gRPC endpoint found by running gcloud compute tpus describe my-tpu:

ipAddress: <XXX>
port: <YYY>

I can easily set up an SSH tunnel to connect to this endpoint from my local machine but I would like to use IAP to create that tunnel instead. I have tried the following:

gcloud compute start-iap-tunnel my-tpu 8470

but I get

 - The resource 'projects/.../zones/.../instances/my-tpu' was not found

This makes sense because a TPU is a not a compute instance, and the command gcloud compute start-iap-tunnel expects an instance name.

Is there any way to use IAP to tunnel to an arbitrary internal IP address? Or more generally, is there any other way that I can use IAP to create a tunnel to my TPU worker?

tpu google-cloud-platform
1 Answers

Yes, it can be done using the internal ip address of the TPU Worker, here is an example:

gcloud alpha compute start-iap-tunnel \
  10.164.0.2 8470 \
  --local-host-port="localhost:$LOCAL_PORT" \
  --region $REGION \
  --network $SUBNET \
  --project $PROJECT

Be aware that Private Google Access must be enabled in the TPU subnet, which can be easily done with the following command:

gcloud compute networks subnets update $SUBNET \
--region=$REGION \
--enable-private-ip-google-access

Just as a reference, here you have an example on how to create a TPU Worker with no external ip address:

gcloud alpha compute tpus tpu-vm create \
  --project $PROJECT \
  --zone $ZONE \
  --internal-ips \
  --version tpu-vm-tf-2.6.0 \
  --accelerator-type v2-8 \
  --network $SUBNET \
   $NAME

AUTHENTICATION

To successfully authenticate the endpoint source of the IAP tunnel, you need to add the SSH keys to the project's metadata following these steps:

  1. Check if you already have SSH keys generated in your endpoint:
ls -1 ~/.ssh/*

#=>

/. . ./id_rsa
/. . ./id_rsa.pub

If you don't have any, you can generate them with the command: ssh-keygen -t rsa -f ~/.ssh/id_rsa -C id_rsa.

  1. Add the SSH keys to your project's metadata:
gcloud compute project-info add-metadata \
--metadata ssh-keys="$(gcloud compute project-info describe \
--format="value(commonInstanceMetadata.items.filter(key:ssh-keys).firstof(value))")
$(whoami):$(cat ~/.ssh/id_rsa.pub)"

#=>

Updated [https://www.googleapis.com/compute/v1/projects/$GCP_PROJECT_NAME].
  1. Assign the iap.tunnelResourceAccessor role to the user:
gcloud projects add-iam-policy-binding $GCP_PROJECT_NAME \
    --member=user:$USER_ID \
    --role=roles/iap.tunnelResourceAccessor
user_169223430
Related