Setup production MLFlow server on google cloud

5 min readJun 5, 2021

Overview

I recently went through the steps to set up an MLFlow server on google cloud with Docker Compose, MySQL, and Google Storage Bucket. It is not the most straightforward thing to do, as the needed knowledge is scattered across many places. Hence, I am hoping to document the steps I took here, in case it may be helpful for you.

Architecture

The overall architecture of the setup can be found in Fig.1 below. MLFlow server includes 3 main components:

The main MLFlow service
The Artifact Store (where all the logged pictures, models, etc. are stored; Normally file systems like S3, GS can be used)
The Backend Store (where all the metadata on the experiment and runs are stored; Normally databases like MySQL, SQLite can be used)

In our architecture, the main MLFlow service and the Backend Store (Implemented using MySQL) are hosted using Docker Compose together on a GCP virtual machine (Ubuntu VM). The Backend Store is deployed as a Google Storage Bucket.

In addition, to make sure the data stored in MySQL won’t be lost when the Backend Store Docker container is deleted, I mounted host volume to save data in a dedicated location on the VM.

Fig. 1. The overall architecture of the setup

Step 1: Install Docker and Docker Compose

As the MLFlow service and the Backend Store will be hosted using Docker containers via Docker Compose, we have to install Docker and Docker Compose as the first step.

Once logged into the Ubuntu VM, run the following steps to install Docker:

$ sudo apt-get remove docker docker-engine docker.io containerd runc$ sudo apt-get update

$ sudo apt-get install \
    apt-transport-https \
    ca-certificates \
    curl \
    gnupg \
    lsb-release$ curl -fsSL https://download.docker.com/linux/ubuntu/gpg | sudo gpg --dearmor -o /usr/share/keyrings/docker-archive-keyring.gpg$ echo \
  "deb [arch=amd64 signed-by=/usr/share/keyrings/docker-archive-keyring.gpg] https://download.docker.com/linux/ubuntu \
  $(lsb_release -cs) stable" | sudo tee /etc/apt/sources.list.d/docker.list > /dev/null$ sudo apt-get update$ sudo apt-get install docker-ce docker-ce-cli containerd.io$ sudo docker run hello-world

After running the last row, you should be able to see the following outputs from the terminal:

Next, we will quickly set up Docker Compose. Simply run the following:

$ sudo curl -L "https://github.com/docker/compose/releases/download/1.29.2/docker-compose-$(uname -s)-$(uname -m)" -o /usr/local/bin/docker-compose$ sudo chmod +x /usr/local/bin/docker-compose

Now, to verify if the Docker Compose has been successfully installed, just type the command docker-compose in the terminal, and it shouldn’t fail if things went successfully.

Step 2: Setup Artifact Store using Google Cloud Storage

To do this, simply create a bucket in Google Cloud Storage, and create a folder will you want to store the MLFlow run artifacts. We will be accessing this Artifact Store using its path defined as:

gs://<bucket name>/<path to your folder>/

Step 3: Setup main MLFlow service, and MySQL Backend Store using Docker Compose

Now we will set up both the main MLflow service and the MySQL Backend Store. They will be hosted by two separate docker containers and managed by Docker Compose together.

For the MLFlow service, we have to write a Dockerfile for building its Docker image. Simply create a Dockerfile and add the following two lines inside:

FROM python:3.6-slim-buster
RUN pip install mlflow google-cloud-storage pymysql

Please note here we will use 3 python packages:

“mlflow” is for running the main MLFlow service
“google-cloud-storage” is needed because our Artifact Store (to be covered later) is using Google Storage Bucket (There’s an equivalent package if you are using AWS S3). Please note that you also need to pip install this on the machine running the ML pipeline that’s using MLFlow server
“pymysql” is to interact with our MySQL Backend Store

For the MySQL Backend Store, we will simply use the MySQL docker image from Docker hub.

Now, we will write the configuration file for Docker Compose to manage the two images. Create a file called “docker-compose.yml” and add the following code inside:

version: ‘3’ services:
    db:
        restart: always
        image: mysql/mysql-server:5.7.28
        container_name: mlflow_db
        expose:
            — “3306”
        environment:
            — MYSQL_DATABASE=${MYSQL_DATABASE}
            — MYSQL_USER=${MYSQL_USER}
            — MYSQL_PASSWORD=${MYSQL_PASSWORD}
        volumes:
            — dbdata:/var/lib/mysql
    web:
        restart: always
        build: ./mlflow   #location of Dockerfile
        image: mlflow_server
        container_name: mlflow_server
        ports:
            — “5000:5000”
        command: mlflow server — backend-store-uri mysql+pymysql://${MYSQL_USER}:${MYSQL_PASSWORD}@db:3306/${MYSQL_DATABASE} — default-artifact-root gs://<bucket name>/<path to your folder>/ — host 0.0.0.0volumes:
   dbdata:

In this file, “db” is the MySQL Backend Store service, and “web” is MLFlow service. Please feel free to read more about the parameters on Docker Compose official documentation.

One more thing to take note of is we are mounting the host volume to the Backend Store MySQL service. So in this way, even if the container is deleted and re-created, the data inside wouldn’t be lost.

We are almost at the end of this setup, however, it wouldn’t work just yet at this very moment. Try building and running the containers via Docker Compose by:

sudo docker-compose up -d — build

You will notice that the MLFlow service keeps failing/restarting. The reason is that to make sure the MLFlow service can access the Backend Store, we have to create a user (with password) and a dedicated database in the MySQL service. To do this, while the MySQL service is still running, firstly get the default root password, and then login as root:

To get the password, run docker logs mysql1 2>&1 | grep GENERATED and copy the password. Then run mysql -u root -p and paste the password to login into MySQL service.

You may be required to reset the password, do it by：

ALTER USER ‘root’@’localhost’ IDENTIFIED BY ‘NewPassword’;

Now, we will create a new user with password, just for the MLFlow to access this MySQL service as Backend Store:

CREATE USER ‘username’ IDENTIFIED BY ‘password’;
CREATE DATABASE ‘mlflowdbname’;
GRANT ALL PRIVILEGES ON mlflowdbname.* TO ‘username’;

We are all set for the MySQL Backend Store service now. Do make sure you set the environment variables so they can be picked up by the MLFlow service to query the MySQL Backend Store:

export MYSQL_DATABASE=mlflowdbname
export MYSQL_USER=username
export MYSQL_PASSWORD=password

One last step, re-start these containers via Docker Compose:

sudo docker-compose stopsudo -E docker-compose up -d — build

Please note that this time I’m using “-E” in docker-compose up command. This is super important, as otherwise the environmental variable couldn’t be picked up by the docker-compose.

VICTORY!!!

Now the MLFlow server is accessible at port 5000. Simply open a browser and type <host IP>:5000 to see the MLFlow UI.