MLOps Pipeline with Python MLflow and Docker
The gap between training a machine learning model in a Jupyter notebook and running it reliably in production is where most ML projects stall. You have a model that performs well on your test set, but deploying it with proper versioning, monitoring, and reproducibility requires an entirely different set of skills and tools. This is where MLOps comes in, and MLflow has emerged as one of the most practical tools for bridging that gap.
I will walk through building a full MLOps pipeline using MLflow and Docker: experiment tracking, model registry, model serving, and containerized deployment.
Why MLflow
MLflow is an open-source platform that addresses four key challenges in the ML lifecycle: tracking experiments, packaging code for reproducibility, managing model versions, and serving models. Unlike heavier platforms that try to own the entire pipeline, MLflow integrates with your existing tools and stays out of the way.
The platform has four core components:
- MLflow Tracking: Log parameters, metrics, and artifacts from your experiments
- MLflow Projects: Package ML code in a reusable, reproducible format
- MLflow Models: A standard format for packaging models for deployment
- MLflow Model Registry: A centralized store for managing model versions and lifecycle stages
Setting Up MLflow with Docker
Let us start by containerizing the MLflow tracking server with a PostgreSQL backend for metadata storage and S3-compatible storage for artifacts.
# docker-compose.yml
version: "3.8"
services:
postgres:
image: postgres:16
environment:
POSTGRES_DB: mlflow
POSTGRES_USER: mlflow
POSTGRES_PASSWORD: mlflow_password
volumes:
- postgres_data:/var/lib/postgresql/data
ports:
- "5432:5432"
minio:
image: minio/minio:latest
command: server /data --console-address ":9001"
environment:
MINIO_ROOT_USER: mlflow_minio
MINIO_ROOT_PASSWORD: mlflow_minio_password
volumes:
- minio_data:/data
ports:
- "9000:9000"
- "9001:9001"
mlflow:
build:
context: .
dockerfile: Dockerfile.mlflow
depends_on:
- postgres
- minio
environment:
MLFLOW_BACKEND_STORE_URI: postgresql://mlflow:mlflow_password@postgres:5432/mlflow
MLFLOW_S3_ENDPOINT_URL: http://minio:9000
AWS_ACCESS_KEY_ID: mlflow_minio
AWS_SECRET_ACCESS_KEY: mlflow_minio_password
ports:
- "5000:5000"
command: >
mlflow server
--backend-store-uri postgresql://mlflow:mlflow_password@postgres:5432/mlflow
--default-artifact-root s3://mlflow-artifacts/
--host 0.0.0.0
--port 5000
volumes:
postgres_data:
minio_data:
# Dockerfile.mlflow
FROM python:3.12-slim
RUN pip install --no-cache-dir \
mlflow==2.15.0 \
psycopg2-binary \
boto3
EXPOSE 5000
Spin up the stack with docker compose up -d. You now have a production-grade MLflow tracking server with persistent storage for both metadata and artifacts.
Experiment Tracking
With the server running, let us train a model and track everything with MLflow. This example uses scikit-learn to build a classification model, but the tracking patterns apply to any framework.
import mlflow
import mlflow.sklearn
from sklearn.ensemble import RandomForestClassifier
from sklearn.datasets import load_wine
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, f1_score, precision_score
import numpy as np
# Point to our tracking server
mlflow.set_tracking_uri("http://localhost:5000")
mlflow.set_experiment("wine-classification")
# Load and split data
data = load_wine()
X_train, X_test, y_train, y_test = train_test_split(
data.data, data.target, test_size=0.2, random_state=42
)
# Define hyperparameter grid
param_grid = [
{"n_estimators": 100, "max_depth": 5, "min_samples_split": 2},
{"n_estimators": 200, "max_depth": 10, "min_samples_split": 5},
{"n_estimators": 300, "max_depth": 15, "min_samples_split": 10},
]
for params in param_grid:
with mlflow.start_run():
# Log parameters
mlflow.log_params(params)
mlflow.log_param("random_state", 42)
# Train model
model = RandomForestClassifier(random_state=42, **params)
model.fit(X_train, y_train)
# Evaluate
predictions = model.predict(X_test)
accuracy = accuracy_score(y_test, predictions)
f1 = f1_score(y_test, predictions, average="weighted")
precision = precision_score(y_test, predictions, average="weighted")
# Log metrics
mlflow.log_metrics({
"accuracy": accuracy,
"f1_score": f1,
"precision": precision
})
# Log feature importance as an artifact
importance = model.feature_importances_
np.savetxt("feature_importance.csv", importance, delimiter=",")
mlflow.log_artifact("feature_importance.csv")
# Log the model itself
mlflow.sklearn.log_model(
model,
artifact_path="model",
registered_model_name="wine-classifier"
)
print(f"Params: {params}")
print(f"Accuracy: {accuracy:.4f}, F1: {f1:.4f}")
Every run captures the parameters, metrics, the trained model, and any additional artifacts. Open http://localhost:5000 in your browser to see the experiment dashboard, compare runs side by side, and visualize how hyperparameters affect performance.
Model Registry and Lifecycle Management
Once you have a model worth deploying, the Model Registry provides version control and lifecycle management.
from mlflow.tracking import MlflowClient
client = MlflowClient(tracking_uri="http://localhost:5000")
# Get the best run from our experiment
experiment = client.get_experiment_by_name("wine-classification")
best_run = client.search_runs(
experiment_ids=[experiment.experiment_id],
order_by=["metrics.accuracy DESC"],
max_results=1
)[0]
print(f"Best run ID: {best_run.info.run_id}")
print(f"Best accuracy: {best_run.data.metrics['accuracy']:.4f}")
# The model was already registered during training via
# registered_model_name. Now manage its lifecycle.
model_name = "wine-classifier"
# Get the latest version
latest_versions = client.get_latest_versions(model_name)
for version in latest_versions:
print(f"Version: {version.version}, Stage: {version.current_stage}")
# Transition the best model to production
client.transition_model_version_stage(
name=model_name,
version=latest_versions[0].version,
stage="Production",
archive_existing_versions=True
)
print(f"Model {model_name} v{latest_versions[0].version} "
f"promoted to Production")
The registry gives you a clear audit trail of which model version is deployed where, who promoted it, and when. This is essential for regulated industries and for debugging production issues.
Serving Models with Docker
MLflow can package any registered model into a Docker container ready for serving. This gives you a REST API endpoint with zero custom code.
# Build a Docker image for the production model
mlflow models build-docker \
--model-uri "models:/wine-classifier/Production" \
--name wine-classifier-serving
# Run the serving container
docker run -p 8080:8080 wine-classifier-serving
Now you can send prediction requests to the model:
import requests
import json
# Sample wine data for prediction
payload = {
"inputs": [
[13.2, 1.78, 2.14, 11.2, 100.0, 2.65, 2.76,
0.26, 1.28, 4.38, 1.05, 3.4, 1050.0]
]
}
response = requests.post(
"http://localhost:8080/invocations",
headers={"Content-Type": "application/json"},
data=json.dumps(payload)
)
print(f"Prediction: {response.json()}")
Custom Serving with FastAPI
For more control over the serving layer — adding authentication, custom preprocessing, or batch predictions — wrap the MLflow model in a FastAPI application.
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import mlflow.pyfunc
import numpy as np
app = FastAPI(title="Wine Classifier API")
# Load the production model at startup
model = mlflow.pyfunc.load_model("models:/wine-classifier/Production")
class PredictionRequest(BaseModel):
features: list[list[float]]
class PredictionResponse(BaseModel):
predictions: list[int]
model_version: str
@app.post("/predict", response_model=PredictionResponse)
async def predict(request: PredictionRequest):
try:
input_data = np.array(request.features)
predictions = model.predict(input_data)
return PredictionResponse(
predictions=predictions.tolist(),
model_version="Production"
)
except Exception as e:
raise HTTPException(status_code=400, detail=str(e))
@app.get("/health")
async def health():
return {"status": "healthy", "model": "wine-classifier"}
# Dockerfile.serving
FROM python:3.12-slim
WORKDIR /app
COPY requirements-serving.txt .
RUN pip install --no-cache-dir -r requirements-serving.txt
COPY serving_app.py .
ENV MLFLOW_TRACKING_URI=http://mlflow:5000
CMD ["uvicorn", "serving_app:app", "--host", "0.0.0.0", "--port", "8080"]
Automated Retraining Pipeline
A complete MLOps pipeline should support automated retraining when new data arrives or model performance degrades.
import mlflow
from mlflow.tracking import MlflowClient
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score
def retrain_pipeline(
new_data_path: str,
performance_threshold: float = 0.90
):
"""Retrain the model if performance drops below threshold."""
client = MlflowClient(tracking_uri="http://localhost:5000")
mlflow.set_tracking_uri("http://localhost:5000")
# Load the current production model
production_model = mlflow.pyfunc.load_model(
"models:/wine-classifier/Production"
)
# Evaluate on new data
X_new, y_new = load_new_data(new_data_path)
current_predictions = production_model.predict(X_new)
current_accuracy = accuracy_score(y_new, current_predictions)
print(f"Current model accuracy on new data: {current_accuracy:.4f}")
if current_accuracy >= performance_threshold:
print("Model performance is acceptable. No retraining needed.")
return
print("Performance below threshold. Retraining...")
mlflow.set_experiment("wine-classification-retrain")
with mlflow.start_run():
# Train new model with combined data
model = RandomForestClassifier(
n_estimators=200, max_depth=10, random_state=42
)
model.fit(X_new, y_new)
new_accuracy = accuracy_score(y_new, model.predict(X_new))
mlflow.log_metric("accuracy", new_accuracy)
if new_accuracy > current_accuracy:
mlflow.sklearn.log_model(
model,
artifact_path="model",
registered_model_name="wine-classifier"
)
print(f"New model registered. "
f"Accuracy improved: {current_accuracy:.4f} -> "
f"{new_accuracy:.4f}")
else:
print("New model did not improve. Keeping current version.")
Key Takeaways
Building an MLOps pipeline does not require expensive, proprietary platforms. MLflow combined with Docker gives you experiment tracking, model versioning, and containerized serving — the core capabilities you need to run ML in production responsibly.
Start with experiment tracking on day one, even before you think you need it. Add the model registry when you have a model worth versioning. Containerize serving when you need reproducible deployments. And automate retraining when you have a feedback loop for model performance. Each layer builds on the last, and MLflow makes it straightforward to adopt incrementally.