TPU Monitoring in MLFlow

Huge thanks to Google TPU Research Cloud for providing me with access to TPU chips! MLFlow sadly doesn’t support TPU monitoring out of the box like with NVIDIA cards and pynvml . Luckily it was very easy to add this feature with the following snippet: def log_tpu_mlflow(step): chip_type, count, data = tpu_system_metrics() for tpu in data: for metric, val in tpu.items(): if metric != "device_id":…