Setting-up MLflow on Google Colab

Question

I frequently use Google Colab to train TF/PyTorch models as Colab provides me with GPU/TPU runtime. Besides, I like working with MLflow to store and compare trained models, tracking progress, sharing, etc. What are the available solutions to use MLflow with Google Colab?

Answer 1

There is a Github issue on this, and although at the time of writing it is still open [EDIT: just closed], contributor dmatrix was kind enough to provide a notebook with a full solution, utilizing pyngrok.

Here is the code (meant to be run on a Colab notebook), reposted here with the implicit permission of the author:

!pip install mlflow --quiet
!pip install pyngrok --quiet

import mlflow

with mlflow.start_run(run_name="MLflow on Colab"):
  mlflow.log_metric("m1", 2.0)
  mlflow.log_param("p1", "mlflow-colab")

# run tracking UI in the background
get_ipython().system_raw("mlflow ui --port 5000 &") # run tracking UI in the background


# create remote tunnel using ngrok.com to allow local port access
# borrowed from https://colab.research.google.com/github/alfozan/MLflow-GBRT-demo/blob/master/MLflow-GBRT-demo.ipynb#scrollTo=4h3bKHMYUIG6

from pyngrok import ngrok

# Terminate open tunnels if exist
ngrok.kill()

# Setting the authtoken (optional)
# Get your authtoken from https://dashboard.ngrok.com/auth
NGROK_AUTH_TOKEN = ""
ngrok.set_auth_token(NGROK_AUTH_TOKEN)

# Open an HTTPs tunnel on port 5000 for http://localhost:5000
ngrok_tunnel = ngrok.connect(addr="5000", proto="http", bind_tls=True)
print("MLflow Tracking UI:", ngrok_tunnel.public_url)

The output of which will be a pyngrok-generated URL like:

MLflow Tracking UI: https://0a23d7a7d0c4.ngrok.io

clicking on which will lead to an MLfLow GUI screen.

(Slight modification of the original code thanks to pyngrok creator, Alex Laird)

Tested with MLflow versions 1.10.0 and 1.11.0.