Connect Colab to paid TPU

Question

I'd like to connect Colab to a PAID TPU (upgrading from the free TPU). I created a JSON key using this guide: https://cloud.google.com/docs/authentication/production#auth-cloud-explicit-python, then uploaded it to Colab. I'm able to connect to my storage but not to the TPU:

%tensorflow_version 2.x
import tensorflow as tf
import os
os.environ['GOOGLE_APPLICATION_CREDENTIALS'] = './gcp-permissions.json'

# Authenticated API request - works.
storage_client = storage.Client.from_service_account_json(
    'gcp-permissions.json')
print(list(storage_client.list_buckets())

#Accessing the TPU - does not work. Request times out.
cluster_resolver = tf.distribute.cluster_resolver.TPUClusterResolver(
    tpu='My-TPU-Name',
    zone='us-central1-a',
    project='My-Project-Name'
)

I've also tried the TPUClusterResolver call with just the tpu name, and with 'credentials=gcp-permissions.json' - same result. I've double-checked that my TPU is up and running in the GCP console. It is not preemptible. What am I missing?

Thanks!

Answer 1

So it looks like you're trying to connect to a paid TPU from your own Google Cloud project from a Colab notebook, is that right? That won't work as the Colab runtime is backed by a GCE VM that is in a different project than your own My-project-name. So instead, you want to also create a GCE VM in that same project and run your training script from that VM. Checkout this tutorial: https://cloud.google.com/tpu/docs/quickstart.