Nothing to TPU in 5 Google Cloud commands

tpu.md

Google Cloud configuration

First: install the CLI program for your distribution: https://cloud.google.com/sdk/install

Parameters

Modify accordingly:

export REGION='us-central1'
export ZONE='us-central1-f'
export PROJECT_NAME='proj'
export PROJECT_ID='proj-id'
export NETWORK='ml-params-net'
export FIREWALL='ml-params-firewall'
export ADDRESSES='ml-params-addresses'
export ACCELERATOR='v2-8'
export VERSION='2.1'
export TPU_NAME='tpu0'
export RANGE='192.168.0.0'
export CIDR='16'
export INSTANCE='ml-params-vm0'

Setup:

gcloud config set compute/region "$REGION"
gcloud config set compute/zone "$ZONE"
gcloud projects create --name "$PROJECT_NAME" --set-as-default "$PROJECT_ID"

Provision

gcloud compute networks create "$NETWORK"
gcloud compute firewall-rules create --network "$NETWORK" --allow='tcp:22,icmp' \
                                     "$FIREWALL"
gcloud compute addresses create --global --purpose='VPC_PEERING' \
                                --addresses="$RANGE" --prefix-length="$CIDR" \
                                --network="$NETWORK" \
                                "$ADDRESSES"
gcloud compute instances create --machine-type='n1-standard-4' --boot-disk-size='500GB' \
                                --image-project='debian-cloud' --image-family='debian-10' \
                                --scopes=cloud-platform --network="$NETWORK" \
                                "$INSTANCE"
gcloud compute tpus create --zone="$ZONE" --description="$TPU_NAME" \
                           --accelerator-type="$ACCELERATOR" --version="$VERSION" \
                           --network="$NETWORK" --range="$RANGE" \
                           "$TPU_NAME"

Test

declare -r TPU_ADDR="$(gcloud compute tpus describe $TPU_NAME --format='value[separator=":"](networkEndpoints.ipAddress, networkEndpoints.port)')"

gcloud compute ssh "$INSTANCE" \
       --command='sudo apt-get install -y python3-venv &&
                  python3 -m venv venv &&
                  . ~/venv/bin/activate &&
                  pip3 install -U pip setuptools wheel &&
                  pip3 install tensorflow tensorflow-datasets jupyter &&
                  curl -sL https://raw.githubusercontent.com/tensorflow/docs/bef6f89/site/en/guide/tpu.ipynb |
                    jupyter nbconvert --to script --stdin --output tpu-tester.py &&
                  sed "s/%/pass \#/g" tpu-tester.py.txt > tpu-tester.py &&
                  COLAB_TPU_ADDR="'"$TPU_ADDR"'" python tpu-tester.py'

(alternate version with guards)

Teardown

gcloud -q compute tpus delete --zone="$ZONE" "$TPU_NAME"
gcloud -q compute instances delete --delete-disks='all' --zone="$ZONE" "$INSTANCE"
gcloud -q compute addresses delete --global "$ADDRESSES"
gcloud -q compute firewall-rules delete "$FIREWALL"
gcloud -q compute networks delete "$NETWORK"