gcloud

Connect to your Cloud TPU using SSH:

$ gcloud compute tpus tpu-vm ssh $TPU_NAME --zone=$ZONE

When you request a slice larger than a single host, Cloud TPU creates a TPU VM for each host. The number of TPU chips per host depends on the TPU version.

To install binaries or run code, connect to each TPU VM using the tpu-vm ssh command.

$ gcloud compute tpus tpu-vm ssh $TPU_NAME --zone=$ZONE

To connect to a specific TPU VM using SSH, use the --worker flag with a 0-based index:

$ gcloud compute tpus tpu-vm ssh $TPU_NAME --zone=$ZONE --worker=1

To run a command on all TPU VMs, use the --worker=all and --command flags:

$ gcloud compute tpus tpu-vm ssh $TPU_NAME \
  --zone=$ZONE \
  --worker=all \
  --command='pip install "jax[tpu]==0.4.20" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'

For Multislice, you can run a command on a single VM using the enumerated TPU name, with each slice prefix and the number appended to it. To run a command on all TPU VMs in all slices, use the --node=all, --worker=all, and --command flags, with an optional --batch-size flag.

$ gcloud compute tpus queued-resources ssh your-queued-resource-id \
  --zone=$ZONE \
  --node=all \
  --worker=all \
  --command='pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html' \
  --batch-size=4

When you connect to VMs using the Google Cloud CLI, Compute Engine creates a persistent SSH key.

Console

To connect to your TPUs in the Google Cloud console, use SSH in your browser:

1. In the Google Cloud console, go to the TPUs page:

   Go to TPUs

2. In the list of TPU VMs, click SSH in the row of the TPU VM that you want to connect to.

   When you connect to TPU VMs using the Google Cloud console, Compute Engine creates an ephemeral SSH key.