Mlflow Mlops Pipeline

---
name: mlflow-mlops-pipeline
metadata:
  category: Machine Learning Operations (MLOps)
description: >-
  Build and manage end-to-end MLOps pipelines using MLflow for experiment tracking, model registry, artifact logging,
  autologging, model evaluation, and deployment serving. Triggers when integrating MLflow tracking servers, configuring
  S3/GCS backend stores, writing custom PyFunc models, automating hyperparameter tuning (Optuna), registering models,
  or building FastAPI/Docker inference services.
compatibility: Python (>= 3.9), MLflow (>= 2.8.0), scikit-learn / PyTorch / TensorFlow, PostgreSQL / S3 / GCS
---

# 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

```text
+-------------------+      +---------------------------------+      +------------------------+
|  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`)

```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.

```python
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:

```python
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.

```python
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`)

```python
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.