mlflow-mlops-pipeline

MLflow MLOps Pipelines & Experiment Tracking

End-to-end MLOps patterns for tracking experiments, managing model registries, versioning artifacts, evaluating metrics, and deploying models using MLflow.

---

1. MLOps Architecture Overview

+-------------------+      +---------------------------------+      +------------------------+
|  Training Data    | ---> | ML Training Pipeline (PyTorch)  | ---> |  MLflow Tracking Server|
|  (S3 / Snowflake) |      | (Autolog + Custom PyFunc)       |      |  (PostgreSQL Backend)  |
+-------------------+      +---------------------------------+      +------------------------+
                                                                                |
                                                                                v
+-------------------+      +---------------------------------+      +------------------------+
|  Inference Client | ---> | FastAPI / MLflow Serve API      | <--- |  MLflow Model Registry |
|  (REST / gRPC)    |      | (Container App / KServe)        |      |  (S3 Artifact Store)   |
+-------------------+      +---------------------------------+      +------------------------+

---

2. Remote Tracking Server & Backend Configuration

Environment Variables setup (.env)

MLFLOW_TRACKING_URI=postgresql://mlflow_user:secure_password@db.prod.internal:5432/mlflow_db
MLFLOW_S3_ENDPOINT_URL=https://s3.us-east-1.amazonaws.com
MLFLOW_PYTHON_BIN=/usr/bin/python3
AWS_ACCESS_KEY_ID=AKIAIOSFODNN7EXAMPLE
AWS_SECRET_ACCESS_KEY=wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY

---

3. Production Training & Tracking Pipeline (train_pipeline.py)

Complete Python implementation showcasing dataset logging, Optuna hyperparameter optimization, autologging, custom metrics, and model registration.

import os
import sys
import logging
import optuna
import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.ensemble import GradientBoostingClassifier
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score, roc_auc_score
import mlflow
import mlflow.sklearn
from mlflow.models.signature import infer_signature

# Configure Logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Initialize MLflow Tracking Server
TRACKING_URI = os.getenv("MLFLOW_TRACKING_URI", "http://localhost:5000")
mlflow.set_tracking_uri(TRACKING_URI)
EXPERIMENT_NAME = "customer_churn_prediction"
mlflow.set_experiment(EXPERIMENT_NAME)

def load_and_preprocess_data():
    """Generates synthetic dataset for illustration."""
    np.random.seed(42)
    X = np.random.randn(1000, 10)
    y = (X[:, 0] + X[:, 1] > 0).astype(int)
    feature_names = [f"feature_{i}" for i in range(10)]
    df = pd.DataFrame(X, columns=feature_names)
    df['target'] = y
    return df, feature_names

def train_and_evaluate():
    df, feature_names = load_and_preprocess_data()
    X = df[feature_names]
    y = df['target']
    
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

    def objective(trial):
        with mlflow.start_run(nested=True):
            n_estimators = trial.suggest_int('n_estimators', 50, 300)
            max_depth = trial.suggest_int('max_depth', 3, 10)
            learning_rate = trial.suggest_float('learning_rate', 0.01, 0.2)

            model = GradientBoostingClassifier(
                n_estimators=n_estimators,
                max_depth=max_depth,
                learning_rate=learning_rate,
                random_state=42
            )
            model.fit(X_train, y_train)
            
            preds = model.predict(X_test)
            acc = accuracy_score(y_test, preds)
            
            # Log params & metrics for sub-run
            mlflow.log_params({
                "n_estimators": n_estimators,
                "max_depth": max_depth,
                "learning_rate": learning_rate
            })
            mlflow.log_metric("val_accuracy", acc)
            
            return acc

    # Parent MLflow Run for Hyperparameter Sweep
    with mlflow.start_run(run_name="optuna_sweep_parent") as parent_run:
        study = optuna.create_study(direction="maximize")
        study.optimize(objective, n_trials=5)

        best_params = study.best_params
        logger.info(f"Best Hyperparameters: {best_params}")

        # Train Final Champion Model
        champion_model = GradientBoostingClassifier(**best_params, random_state=42)
        champion_model.fit(X_train, y_train)

        # Evaluate Champion Model
        y_pred = champion_model.predict(X_test)
        y_proba = champion_model.predict_proba(X_test)[:, 1]

        metrics = {
            "accuracy": accuracy_score(y_test, y_pred),
            "precision": precision_score(y_test, y_pred),
            "recall": recall_score(y_test, y_pred),
            "f1_score": f1_score(y_test, y_pred),
            "roc_auc": roc_auc_score(y_test, y_proba)
        }

        # Log metrics to parent run
        mlflow.log_params(best_params)
        mlflow.log_metrics(metrics)

        # Log Model Signature & Input Example
        signature = infer_signature(X_test, champion_model.predict(X_test))
        input_example = X_test.head(3)

        # Register Champion Model to MLflow Model Registry
        model_info = mlflow.sklearn.log_model(
            sk_model=champion_model,
            artifact_path="model",
            signature=signature,
            input_example=input_example,
            registered_model_name="CustomerChurnClassifier"
        )
        logger.info(f"Model logged successfully: {model_info.model_uri}")

if __name__ == "__main__":
    train_and_evaluate()

---

4. Custom PyFunc Wrapper for Preprocessing Pipelines (custom_pyfunc.py)

When your inference workflow requires custom pre-processing or feature transformations before executing model prediction:

import mlflow.pyfunc
import numpy as np
import pandas as pd

class ChurnPipelineWrapper(mlflow.pyfunc.PythonModel):
    def __init__(self, model, scaler):
        self.model = model
        self.scaler = scaler

    def predict(self, context, model_input: pd.DataFrame) -> np.ndarray:
        """Custom inference method handling pre-processing."""
        # 1. Fill missing values
        cleaned_input = model_input.fillna(0)
        
        # 2. Scale features
        scaled_input = self.scaler.transform(cleaned_input)
        
        # 3. Model predict probabilities
        probabilities = self.model.predict_proba(scaled_input)[:, 1]
        
        return np.where(probabilities > 0.5, 1, 0)

# Save custom PyFunc model
# mlflow.pyfunc.log_model(
#     artifact_path="churn_pyfunc_model",
#     python_model=ChurnPipelineWrapper(model, scaler),
#     registered_model_name="CustomerChurnPyFunc"
# )

---

5. Automated Model Registry & Production Promotion (promote_model.py)

Script for CI/CD pipelines to evaluate candidate models and transition stages dynamically.

import os
import logging
from mlflow.tracking import MlflowClient

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

client = MlflowClient(tracking_uri=os.getenv("MLFLOW_TRACKING_URI"))
MODEL_NAME = "CustomerChurnClassifier"

def promote_latest_model():
    # Fetch latest version of model
    latest_versions = client.get_latest_versions(MODEL_NAME, stages=["None"])
    if not latest_versions:
        logger.info("No new model versions found for promotion.")
        return

    latest_version = latest_versions[0].version
    run_id = latest_versions[0].run_id

    # Fetch metric performance
    run_data = client.get_run(run_id)
    accuracy = run_data.data.metrics.get("accuracy", 0.0)

    ACCURACY_THRESHOLD = 0.85

    if accuracy >= ACCURACY_THRESHOLD:
        logger.info(f"Model version {latest_version} passed check (Accuracy: {accuracy:.4f}). Transitioning to Production.")
        
        # Archive current Production models
        for mv in client.search_model_versions(f"name='{MODEL_NAME}'"):
            if mv.current_stage == "Production":
                client.transition_model_version_stage(
                    name=MODEL_NAME,
                    version=mv.version,
                    stage="Archived"
                )

        # Promote new model to Production
        client.transition_model_version_stage(
            name=MODEL_NAME,
            version=latest_version,
            stage="Production"
        )
    else:
        logger.warning(f"Model version {latest_version} failed validation (Accuracy: {accuracy:.4f} < {ACCURACY_THRESHOLD}).")

if __name__ == "__main__":
    promote_latest_model()

---

6. Serving MLflow Models via FastAPI (serve_fastapi.py)

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import pandas as pd
import mlflow.pyfunc
import os

app = FastAPI(title="MLflow Model Serving API", version="1.0.0")

# Load model from MLflow Registry on startup
MODEL_URI = os.getenv("MODEL_URI", "models:/CustomerChurnClassifier/Production")
model = None

@app.on_event("startup")
def load_model():
    global model
    mlflow.set_tracking_uri(os.getenv("MLFLOW_TRACKING_URI", "http://localhost:5000"))
    model = mlflow.pyfunc.load_model(MODEL_URI)

class InferenceRequest(BaseModel):
    features: list[dict]

@app.post("/predict")
def predict(request: InferenceRequest):
    if not model:
        raise HTTPException(status_code=500, detail="Model not initialized")
    
    df = pd.DataFrame(request.features)
    predictions = model.predict(df)
    return {"predictions": predictions.tolist()}

---

7. MLOps Best Practices Checklist

1. Artifact Store Isolation: Store artifacts in dedicated cloud bucket prefixes per environment (s3://my-ml-artifacts-prod/).
2. Model Signatures: Always log explicit inputs and outputs signatures using mlflow.models.signature.infer_signature.
3. Environment Lock: Log conda.yaml or requirements.txt environment locks alongside models to guarantee reproducible containerized deployments.
4. Deterministic Run Names: Set clear, human-readable run names (mlflow.start_run(run_name="xgboost_baseline_v1")) instead of relying on default UUID strings.