Running cloud TPU profiler in Google Colab environment

Question

I am running a Google Colab notebook and am trying to capture TPU profiling data for use in TensorBoard, however I can't get capture_tpu_profile to run in the background while running my TensorFlow code.

So far I tried to run the capture process in the background with:

!capture_tpu_profile --logdir=gs://<my_logdir> --tpu=$COLAB_TPU_ADDR &

and

!bg capture_tpu_profile --logdir=gs://<my_logdir> --tpu=$COLAB_TPU_ADDR

Answer 1

Turns out a way to do this is to start the process from python directly like this (I also had to modify the parameter from --tpu to --service_addr):

import subprocess
subprocess.Popen(["capture_tpu_profile","--logdir=gs://<my_logdir>", "--service_addr={}".format(os.environ['COLAB_TPU_ADDR'])])

the check=True makes the command raise an Exception if it fails.

Answer 2

One way to do this is to use the TPUProfilerHook

https://github.com/tensorflow/tpu/blob/master/models/common/tpu_profiler_hook.py

Which runs the profiler as a session hook.

Example here https://github.com/tensorflow/tpu/blob/5d838047af0163bdf7b97b9404648dc2961c4b63/models/official/resnet/resnet_main.py#L699