TPU Monitoring in MLFlow Part 2

Huge thanks to Google TPU Research Cloud for providing me with access to TPU chips! As promised, this is a follow up post with a cleaner way to track Google’s TPU performance in MLFlow while training your models. This will spawn a background process (not thread, jaxlib didn’t like that) that will periodically send to the MLFlow Tracking server the relevant TPU information. import multiprocessing…