Switching software versions on your Cloud TPU

Overview

Using Cloud TPU requires selecting a software version to match the software running on your local VM. This software version can now be switched on a running Cloud TPU, without deleting and recreating the TPU. This also enables configuring the Cloud TPU with specific nightly versions of software frameworks. It is still recommended to select a supported version of these frameworks.

Usage

The recommended way to switch versions is to use the cloud-tpu-client python library.

Example usage for TensorFlow.

import tensorflow as tf
from cloud_tpu_client import Client

c = Client()
c.configure_tpu_version(tf.__version__, restart_type='ifNeeded')

This will configure the Cloud TPU to match the exact TensorFlow version running on your local VM, this includes official releases as well as dated nightlies.

The library communicates directly with the Cloud TPU so this code needs to be run in a VM in the same network. It's recommended to run this within the code for the rest of your model.

Additional software options

Tensorflow includes a tf.__version__ string which is the simplest way to configure the correct version. Other software options include:

  • PyTorch - pytorch-0.1, pytorch-0.5-dev20200320, pytorch-nightly
  • Jax - tpu_driver, tpu_driver0.1-dev20200320, tpu_driver_nightly

For example to configure a TPU to run with the latest nightly build of PyTorch.

from cloud_tpu_client import Client
c = Client()
c.configure_tpu_version('pytorch-nightly', restart_type='ifNeeded')