Posted inAI

Deploy MLflow on Docker: Track ML Experiments, Manage Model Registry in Real MLOps Pipelines

AI tutorial - IT technology blog
AI tutorial - IT technology blog

The 2 AM Problem: Which Model Was That Again?

It was 2:17 AM when the message came in — production accuracy had dropped significantly after the latest model deployment. My first instinct was to roll back, but I had a bigger problem: I had no idea which exact model version was running, what hyperparameters trained it, or which dataset split produced it. I had run dozens of experiments across three machines over two weeks, naming files things like model_final.pkl, model_final_v2.pkl, and the classic model_ACTUALLY_final.pkl.

That night cost us four hours of scrambling through Jupyter notebooks, git logs, and CSV files to reconstruct what should have been a five-minute lookup. If you’ve shipped a machine learning project into production without proper experiment tracking, you already know this feeling. That incident was what pushed me to properly set up MLflow — and in my real-world experience, this is one of the essential skills to master before you touch anything that resembles a production ML system.

What MLflow Actually Does Before We Touch Docker

MLflow is an open-source platform for managing the end-to-end ML lifecycle. It’s not magic — it’s structured bookkeeping for your models and experiments. Four components make up the platform:

  • Tracking Server — logs parameters, metrics, and artifacts (model files, plots) for each training run
  • Model Registry — a central store to version, stage, and promote models (Staging → Production → Archived)
  • MLflow Models — a standard packaging format that works with scikit-learn, PyTorch, TensorFlow, and others
  • Projects — reproducible packaging of your training code (less commonly used, but useful for teams)

For most teams, the Tracking Server and Model Registry alone solve 90% of the chaos. That’s what we’ll deploy today — with Docker Compose so the setup is reproducible and doesn’t rot on someone’s laptop.

Setting Up MLflow with Docker Compose

We’ll run MLflow backed by a PostgreSQL database (so run metadata survives container restarts) and MinIO as an S3-compatible artifact store (so model files don’t live inside the container filesystem).

Project Structure

mlflow-docker/
├── docker-compose.yml
├── .env
└── mlflow/
    └── Dockerfile

The .env File

# .env
POSTGRES_USER=mlflow
POSTGRES_PASSWORD=mlflow_secret
POSTGRES_DB=mlflow
MINIO_ROOT_USER=minioadmin
MINIO_ROOT_PASSWORD=minioadmin
MINIO_BUCKET=mlflow-artifacts
MLFLOW_S3_ENDPOINT_URL=http://minio:9000
AWS_ACCESS_KEY_ID=minioadmin
AWS_SECRET_ACCESS_KEY=minioadmin

docker-compose.yml

version: "3.8"

services:
  postgres:
    image: postgres:15
    environment:
      POSTGRES_USER: ${POSTGRES_USER}
      POSTGRES_PASSWORD: ${POSTGRES_PASSWORD}
      POSTGRES_DB: ${POSTGRES_DB}
    volumes:
      - postgres_data:/var/lib/postgresql/data
    healthcheck:
      test: ["CMD", "pg_isready", "-U", "${POSTGRES_USER}"]
      interval: 10s
      retries: 5

  minio:
    image: minio/minio:latest
    command: server /data --console-address ":9001"
    environment:
      MINIO_ROOT_USER: ${MINIO_ROOT_USER}
      MINIO_ROOT_PASSWORD: ${MINIO_ROOT_PASSWORD}
    ports:
      - "9000:9000"
      - "9001:9001"
    volumes:
      - minio_data:/data

  minio-init:
    image: minio/mc:latest
    depends_on:
      - minio
    entrypoint: >
      /bin/sh -c "
      sleep 5;
      mc alias set local http://minio:9000 ${MINIO_ROOT_USER} ${MINIO_ROOT_PASSWORD};
      mc mb local/${MINIO_BUCKET} --ignore-existing;
      exit 0;
      "
    environment:
      MINIO_ROOT_USER: ${MINIO_ROOT_USER}
      MINIO_ROOT_PASSWORD: ${MINIO_ROOT_PASSWORD}
      MINIO_BUCKET: ${MINIO_BUCKET}

  mlflow:
    image: ghcr.io/mlflow/mlflow:v2.13.0
    ports:
      - "5000:5000"
    environment:
      MLFLOW_S3_ENDPOINT_URL: ${MLFLOW_S3_ENDPOINT_URL}
      AWS_ACCESS_KEY_ID: ${AWS_ACCESS_KEY_ID}
      AWS_SECRET_ACCESS_KEY: ${AWS_SECRET_ACCESS_KEY}
    command: >
      mlflow server
      --backend-store-uri postgresql://${POSTGRES_USER}:${POSTGRES_PASSWORD}@postgres/${POSTGRES_DB}
      --default-artifact-root s3://${MINIO_BUCKET}/
      --host 0.0.0.0
      --port 5000
    depends_on:
      postgres:
        condition: service_healthy
      minio-init:
        condition: service_completed_successfully

volumes:
  postgres_data:
  minio_data:

Start the Stack

docker compose up -d

# Check everything is healthy
docker compose ps

# MLflow UI → http://localhost:5000
# MinIO Console → http://localhost:9001

Give it about 30 seconds on first startup while PostgreSQL initializes and MinIO creates the bucket. Once mlflow container shows healthy, the UI is live at port 5000.

Logging Your First Experiment

On your local machine (or wherever your training code runs), install the MLflow client and boto3 for S3 artifact uploads:

pip install mlflow boto3

Basic Experiment Tracking

import mlflow
import mlflow.sklearn
from sklearn.ensemble import RandomForestClassifier
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score
import os

os.environ["MLFLOW_S3_ENDPOINT_URL"] = "http://localhost:9000"
os.environ["AWS_ACCESS_KEY_ID"] = "minioadmin"
os.environ["AWS_SECRET_ACCESS_KEY"] = "minioadmin"

mlflow.set_tracking_uri("http://localhost:5000")
mlflow.set_experiment("iris-classification")

X, y = load_iris(return_X_y=True)
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

for n_estimators in [50, 100, 200]:
    with mlflow.start_run(run_name=f"rf-{n_estimators}-trees"):
        model = RandomForestClassifier(n_estimators=n_estimators, random_state=42)
        model.fit(X_train, y_train)

        accuracy = accuracy_score(y_test, model.predict(X_test))

        # Log parameters and metrics
        mlflow.log_param("n_estimators", n_estimators)
        mlflow.log_metric("accuracy", accuracy)

        # Log the model itself as an artifact
        mlflow.sklearn.log_model(
            sk_model=model,
            artifact_path="model",
            registered_model_name="iris-random-forest"
        )

        print(f"n_estimators={n_estimators} → accuracy={accuracy:.4f}")

Run this script, then open http://localhost:5000. You’ll see three runs under the iris-classification experiment, each with its parameters, metrics, and model artifacts stored in MinIO. Click any run to see the full lineage — dataset info, code version, environment, everything you’d need to reproduce it.

Managing the Model Registry

The three runs above already registered the model automatically via registered_model_name. Now let’s manage versions through stages — this is what makes rollbacks a 30-second task instead of a 4-hour nightmare.

Promoting Models Through Stages

from mlflow.tracking import MlflowClient

client = MlflowClient(tracking_uri="http://localhost:5000")

# List all versions of our model
versions = client.search_model_versions("name='iris-random-forest'")
for v in versions:
    print(f"Version {v.version} — run_id={v.run_id} — stage={v.current_stage}")

# Promote the best-performing version to Staging
client.transition_model_version_stage(
    name="iris-random-forest",
    version="3",  # the 200-tree version
    stage="Staging"
)

# After validation passes, move to Production
client.transition_model_version_stage(
    name="iris-random-forest",
    version="3",
    stage="Production"
)

# Archive the old production model
client.transition_model_version_stage(
    name="iris-random-forest",
    version="1",
    stage="Archived"
)

Loading a Production Model for Inference

import mlflow.sklearn
import os

os.environ["MLFLOW_S3_ENDPOINT_URL"] = "http://localhost:9000"
os.environ["AWS_ACCESS_KEY_ID"] = "minioadmin"
os.environ["AWS_SECRET_ACCESS_KEY"] = "minioadmin"

mlflow.set_tracking_uri("http://localhost:5000")

# Always loads whatever is currently tagged Production
model = mlflow.sklearn.load_model("models:/iris-random-forest/Production")

sample = [[5.1, 3.5, 1.4, 0.2]]
print(model.predict(sample))  # [0]

The critical part here is models:/iris-random-forest/Production — your inference code never hardcodes a version number. When you promote version 4 to Production, the next prediction call picks it up automatically. Rolling back is equally simple: transition an older version back to Production.

What You Now Have

At this point, the stack gives you a fully self-hosted MLOps backbone: experiment metadata in PostgreSQL (survives restarts, queryable), model artifacts in MinIO (S3-compatible, scalable), and a clean Model Registry with promotion stages that map directly to your deployment pipeline. Every training run is auditable — who trained it, what code, what data split, what metrics.

The next time something breaks at 2 AM, you open the MLflow UI, filter by the model name, check which version is in Production, compare its training metrics against the previous version, and roll back in under two minutes. That’s the difference between a recoverable incident and a multi-hour fire.

For teams, the natural next step is wiring your CI/CD pipeline to auto-register models after successful training runs and gate promotion to Production behind metric thresholds — but the foundation you’ve built here handles that cleanly. The infrastructure is in place; the workflow policies are yours to define.

Share:
Tags: