Use Identity Aware Proxy to tunnel to a TPU - google-cloud-platform

Question

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?

Answer 1

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:
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.
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].
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