How to upgrade from Colab TPU to paid TPU?

Question

I've been using Colab with a TPU over the past few days and it's been great. Now, I'd like to upgrade to a paid TPU on GCE. I've been following the quickstart here: https://cloud.google.com/tpu/docs/quickstart

I've created a storage bucket and a TPU using tpu up. Now would like to connect to the TPU using my notebook, hosted with GCP AI Notebooks (https://console.cloud.google.com/ai-platform/notebooks). When I was on Colab, I could access it using:

cluster_resolver = tf.distribute.cluster_resolver.TPUClusterResolver()
tf.config.experimental_connect_to_cluster(cluster_resolver)

When I run the same line in a GCP AI Notebook...

cluster_resolver = tf.distribute.cluster_resolver.TPUClusterResolver('MyTPUName')
tf.config.experimental_connect_to_cluster(cluster_resolver)

...it throws this error

HttpError: <HttpError 403 when requesting https://tpu.googleapis.com/v1/projects/local-dialect-[MyProjectID]/locations/us-east1-c/nodes/[MyTPUName]?alt=json returned "Location us-east1-c is not found or access is unauthorized.">

How do I fix it? BTW, if anyone from Google is reading this: If you're going to charge $1-5/hr for this product, how about an "upgrade" button from Colab? Everybody wins.

Thank you!

Answer 1

Turns out I needed a few extra parameters to get through security. This worked:

cluster_resolver = tf.distribute.cluster_resolver.TPUClusterResolver( \
        tpu='[my tpu name]', 
        zone='us-central1-a', 
        project='[my project name]')
tf.config.experimental_connect_to_cluster(cluster_resolver)