MLflow with remote Tracking Server, backend and artifact stores

MLflow is an open source platform to manage and track the ML lifecycle, including experimentation, reproducibility, deployment, and a central model registry.

MLflow has the ability to work with distributed architectures, which allows the tracking server, backend store, and artifact store to be located on remote hosts. For instance, consider the following scenario where a remote MLflow Tracking Server is used, along with a Postgres database for backend entity storage, and an S3 bucket for artifact storage.

Why use Docker with MLflow?

By containerizing the MLflow server, it becomes easier to deploy and manage the server across different environments. Additionally, by using a remote artifact store and backend store, it is possible to use MLflow in a distributed architecture. This can be particularly useful when working with large datasets or when collaborating with multiple team members on a project.

Postgresql Side:

first thing to do is to create a separate user and database that dedicated to mlflow, by creating a separate user and database for each service, you can limit the scope of access that each service has to the database. If a service has a bug or vulnerability that allows an attacker to access the database, the attacker will only be able to access the data associated with that service, rather than the entire database. In addition, having a separate database for each service can make it easier to troubleshoot errors. If an error occurs, you can look at the logs and data associated with the specific service, rather than having to sift through a large database with data from multiple services.

to create a new user in PostgreSQL you should open a terminal where PostgreSQL is running and log in as a superuser :

sudo -u postgres psql

Create the new user “MLFLOW_USER” and set the password “PASSWORD”

CREATE USER <<MLFLOW_USER>> WITH PASSWORD <<'PASSWORD'>>;

Grant the user permission to create a database:

ALTER USER <<MLFLOW_USER>> CREATEDB;

create database “mlflow_db” that will be used by MLflow:

CREATE DATABASE <<MLFLOW_DB>> OWNER <<MLFLOW_USER>>;

MinIO S3:

MLflow major ability is the ability to log and store the artifacts generated during the machine learning workflow, such as models, data, and visualizations. These artifacts can be critical for reproducing and understanding the results of a machine learning experiment. To manage and store those artifacts in a scalable, durable, and secure way, you need to use S3 bucket.

To do that open MinIO web interface http://<IP_ADRESSE>:9000 in a web browser. After log in with your credentials, click on “Bucket” and afterwards “Create Bucket” button.

In the “Bucket Name” field, enter a unique name for your bucket, something like “mlflow”. (Bucket names must be globally unique across all of S3)

• (Optional) Set a versioning policy for the bucket to enable versioning for objects.
• (Optional) Set a Quota for the bucket to to limit the amount of data in the bucket.

MLflow Docker Image:

Next is to create a Dockerfile that builds a Docker image with MLflow and use S3 minIO Bucket as an artifact store und PostgreSQL as a backend store.

Dockerfile

# Start from a base image with Python installed
FROM python:3.9
# Install mlflow and dependencies
RUN pip install \
    mlflow \
    psycopg2 \
    boto3 && \
    mkdir /mlflow/
# Expose the default MLflow server port
EXPOSE 5000
# Start the MLflow server when the container starts
CMD mlflow server \
    --host 0.0.0.0 \
    --port 5000 \
    --artifacts-destination s3://mlflow \
    --backend-store-uri postgresql://<<MLFLOW_USER>>:<<PASSWORD>>!@<<IP_ADDRESS>>:5432/<<MLFLOW_DB>>

• --backend-store-uri specifies the URI for the MLflow backend store. The backend store is responsible for storing experiment and run metadata, such as run parameters, metrics, and artifacts. In this case, we're setting the backend store URI to the value to mlflow database that we did created at the beginning.
• --artifacts-destination specifies the destination URI for storing artifacts. Artifacts are any files or data that are produced or used during the course of a run, such as trained models, visualizations, or data files. In this case, we're setting the artifacts-destination to the same value as the S3 bucket that’s was before created.
• --host specifies the network interface to bind to. In this case, we're setting it to 0.0.0.0, which means that the MLflow server will be accessible from any network interface on the container host.

Docker build

now is the turn of building the Docker image using the Docker file, to do that, let’s run the docker build command to build the Docker image. This command uses the Dockerfile to build the image and saves it the local machine.

sudo docker build -t mlflow-server:latest .

Push Docker image in the private Registry

to be able to push the new created image to the registry when need to log in, to do that we will run the following command

docker login -u <user> -p <password or acces-token > my_registry.com
docker push mlflow-server:<TAG>

Run MLflow Container

To run MLflow as a container and to use a remote S3 bucket as the artifact store, you will need to configure your MLflow container to use the appropriate credentials and endpoint to access the bucket.

so first thing to do is to create an environment variable file to configure all the needed settings, the file will be named mlflow.env and will look like the following:

AWS_ACCESS_KEY_ID=<<YOUR_AWS_ACCESS_KEY_ID>>
AWS_SECRET_ACCESS_KEY=<<AWS_SECRET_ACCESS_KEY>>
S3_ENDPOINT_URL=http://<<IP_ADDRESS>>:9000
MLFLOW_S3_IGNORE_TLS=true

📌 By S3_ENDPOINT_URL You must be sure that you are given the minIO API port number and not the console port number

next thing is the real a shell script that will make easier to run the container each time needed, the shell script name will be ‘mlflow.sh’ and will look like the following:

NoneBashCSSCC#ElixirGoHTMLObjective-CJavaJavaScriptJSONKotlinMermaid DiagramPerlPHPPowershellPythonRubyRustSQLSoliditySwiftTypeScriptYAMLCopy

docker run --rm \
         -p 5000:5000 \
         --name=mlflow-server \
         --env-file=mlflow.env \
        -d  mlflow-server:<<TAG>>

make the file executable

chmod +x mlflow.sh

and execute the shell script with the command ./mlflow.sh

after starting your container you will be able to access mlflow:

Lastly, I warmly invite you to learn more about our company. Visit our website to explore our services in depth.

Whether you have inquiries or seek more information, our support chat is available, or you can reach out through various communication options provided on our website.

We are eager to assist and look forward to connecting with you.