Prometheus + Grafana for ML Model Monitoring (with Airflow Retraining)

python

from fastapi import FastAPI
from prometheus_client import Counter, Histogram, Gauge, generate_latest
from prometheus_fastapi_instrumentator import Instrumentator
import time
import numpy as np

app = FastAPI(title="Fraud Detection API")

# Define custom metrics
predictions_total = Counter(
    'ml_predictions_total',
    'Total number of predictions made',
    ['model_version', 'prediction']
)

prediction_latency = Histogram(
    'ml_prediction_latency_seconds',
    'Prediction latency in seconds',
    buckets=[0.01, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0]
)

prediction_errors = Counter(
    'ml_prediction_errors_total',
    'Total prediction errors',
    ['error_type']
)

fraud_rate = Gauge(
    'ml_fraud_rate_current',
    'Current fraud detection rate (rolling window)'
)

model_accuracy = Gauge(
    'ml_model_accuracy',
    'Model accuracy (when ground truth available)',
    ['model_version']
)

feature_stats = Gauge(
    'ml_feature_value',
    'Feature statistics',
    ['feature_name', 'stat_type']
)

# Add default FastAPI metrics (request count, latency, etc.)
Instrumentator().instrument(app).expose(app)

# Load model
MODEL_VERSION = "v1.2.0"
model = load_model()  # Your model loading logic

@app.post("/predict")
async def predict(transaction: dict):
    """Make fraud prediction with full instrumentation"""
    start_time = time.time()
    
    try:
        # Extract features
        features = extract_features(transaction)
        
        # Track feature statistics
        for feature_name, value in features.items():
            feature_stats.labels(
                feature_name=feature_name,
                stat_type='current'
            ).set(value)
        
        # Make prediction
        prediction = model.predict([list(features.values())])[0]
        prediction_proba = model.predict_proba([list(features.values())])[0]
        
        # Record metrics
        latency = time.time() - start_time
        prediction_latency.observe(latency)
        
        predictions_total.labels(
            model_version=MODEL_VERSION,
            prediction='fraud' if prediction == 1 else 'legitimate'
        ).inc()
        
        # Track fraud rate (you'd use a rolling window in production)
        if prediction == 1:
            fraud_rate.set(fraud_rate._value.get() + 0.01)
        
        return {
            "prediction": int(prediction),
            "probability": float(prediction_proba[1]),
            "model_version": MODEL_VERSION,
            "latency_ms": round(latency * 1000, 2)
        }
        
    except Exception as e:
        # Record error
        prediction_errors.labels(error_type=type(e).__name__).inc()
        raise e

@app.post("/feedback")
async def feedback(transaction_id: str, actual_fraud: bool):
    """Receive ground truth for accuracy tracking"""
    # In production, you'd look up the prediction from a database
    predicted_fraud = get_prediction(transaction_id)
    
    # Calculate and update accuracy
    is_correct = (predicted_fraud == actual_fraud)
    
    # Update accuracy metric (simplified - use proper calculation in prod)
    current_accuracy = model_accuracy.labels(
        model_version=MODEL_VERSION
    )._value.get()
    
    # Exponential moving average
    new_accuracy = 0.9 * current_accuracy + 0.1 * (1.0 if is_correct else 0.0)
    model_accuracy.labels(model_version=MODEL_VERSION).set(new_accuracy)
    
    return {"status": "feedback recorded"}

@app.get("/health")
async def health():
    """Health check endpoint"""
    return {
        "status": "healthy",
        "model_version": MODEL_VERSION
    }

@app.get("/metrics")
async def metrics():
    """Prometheus metrics endpoint"""
    from starlette.responses import Response
    return Response(content=generate_latest(), media_type="text/plain")

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)

python

from airflow import DAG
from airflow.operators.python import PythonOperator
from airflow.sensors.python import PythonSensor
from datetime import datetime, timedelta
import requests
import pandas as pd
from sklearn.ensemble import RandomForestClassifier
import joblib
import mlflow

# Default args
default_args = {
    'owner': 'ml-team',
    'depends_on_past': False,
    'email_on_failure': True,
    'email': ['alerts@company.com'],
    'retries': 2,
    'retry_delay': timedelta(minutes=5)
}

# DAG definition
dag = DAG(
    'ml_model_retraining',
    default_args=default_args,
    description='Automated ML model retraining based on performance',
    schedule_interval='0 2 * * *',  # Daily at 2 AM
    start_date=datetime(2025, 1, 1),
    catchup=False,
    tags=['ml', 'retraining']
)

def check_model_performance():
    """Query Prometheus to check if retraining is needed"""
    prometheus_url = "http://prometheus:9090/api/v1/query"
    
    # Check accuracy
    accuracy_query = "ml_model_accuracy"
    response = requests.get(prometheus_url, params={'query': accuracy_query})
    accuracy = float(response.json()['data']['result'][0]['value'][1])
    
    # Check P95 latency
    latency_query = "histogram_quantile(0.95, rate(ml_prediction_latency_seconds_bucket[1h]))"
    response = requests.get(prometheus_url, params={'query': latency_query})
    p95_latency = float(response.json()['data']['result'][0]['value'][1])
    
    print(f"Current accuracy: {accuracy}")
    print(f"Current P95 latency: {p95_latency}s")
    
    # Decide if retraining is needed
    needs_retraining = (
        accuracy < 0.90 or  # Accuracy dropped below 90%
        p95_latency > 0.3   # Latency above 300ms
    )
    
    if needs_retraining:
        print("⚠️ Model needs retraining!")
    else:
        print("✅ Model performance is acceptable")
    
    return needs_retraining

def prepare_training_data(**context):
    """Pull fresh training data"""
    print("📊 Preparing training data...")
    
    # In production, query your database
    # For this example, we'll simulate it
    query = """
        SELECT * FROM transactions 
        WHERE timestamp > NOW() - INTERVAL '90 days'
        AND label IS NOT NULL
    """
    
    # df = pd.read_sql(query, connection)
    df = load_recent_data()  # Your data loading logic
    
    print(f"Loaded {len(df)} training samples")
    
    # Save to temp location
    df.to_parquet('/tmp/training_data.parquet')
    
    return len(df)

def train_new_model(**context):
    """Train new model with fresh data"""
    print("🔧 Training new model...")
    
    # Load data
    df = pd.read_parquet('/tmp/training_data.parquet')
    
    X = df.drop(['label', 'transaction_id', 'timestamp'], axis=1)
    y = df['label']
    
    from sklearn.model_selection import train_test_split
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.2, random_state=42
    )
    
    # Start MLflow run
    with mlflow.start_run(run_name=f"retrain_{datetime.now().strftime('%Y%m%d_%H%M')}"):
        
        # Train model
        model = RandomForestClassifier(
            n_estimators=100,
            max_depth=10,
            min_samples_split=5,
            random_state=42,
            n_jobs=-1
        )
        
        model.fit(X_train, y_train)
        
        # Evaluate
        from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score
        
        y_pred = model.predict(X_test)
        
        accuracy = accuracy_score(y_test, y_pred)
        precision = precision_score(y_test, y_pred)
        recall = recall_score(y_test, y_pred)
        f1 = f1_score(y_test, y_pred)
        
        print(f"New model accuracy: {accuracy:.4f}")
        print(f"New model precision: {precision:.4f}")
        print(f"New model recall: {recall:.4f}")
        print(f"New model F1: {f1:.4f}")
        
        # Log to MLflow
        mlflow.log_param("n_estimators", 100)
        mlflow.log_param("max_depth", 10)
        mlflow.log_param("training_samples", len(X_train))
        
        mlflow.log_metric("accuracy", accuracy)
        mlflow.log_metric("precision", precision)
        mlflow.log_metric("recall", recall)
        mlflow.log_metric("f1", f1)
        
        # Save model
        model_path = f'/models/fraud_model_{datetime.now().strftime("%Y%m%d_%H%M")}.pkl'
        joblib.dump(model, model_path)
        
        mlflow.sklearn.log_model(model, "model")
        
        # Push to XCom for next task
        context['task_instance'].xcom_push(key='new_model_path', value=model_path)
        context['task_instance'].xcom_push(key='new_accuracy', value=accuracy)
        
        return accuracy

def validate_new_model(**context):
    """Validate new model before deployment"""
    print("✅ Validating new model...")
    
    # Get current production metrics
    prometheus_url = "http://prometheus:9090/api/v1/query"
    response = requests.get(prometheus_url, params={'query': 'ml_model_accuracy'})
    current_accuracy = float(response.json()['data']['result'][0]['value'][1])
    
    # Get new model accuracy from XCom
    new_accuracy = context['task_instance'].xcom_pull(
        task_ids='train_model',
        key='new_accuracy'
    )
    
    print(f"Current production accuracy: {current_accuracy:.4f}")
    print(f"New model accuracy: {new_accuracy:.4f}")
    
    # Validate
    improvement = new_accuracy - current_accuracy
    
    if new_accuracy >= current_accuracy * 0.98:  # At least 98% of current
        print(f"✅ New model approved! Improvement: {improvement:+.4f}")
        return True
    else:
        print(f"❌ New model rejected. Performance drop: {improvement:.4f}")
        raise ValueError("New model performs worse than current model")

def deploy_model(**context):
    """Deploy new model to production"""
    print("🚀 Deploying new model...")
    
    model_path = context['task_instance'].xcom_pull(
        task_ids='train_model',
        key='new_model_path'
    )
    
    # Copy to production location
    import shutil
    shutil.copy(model_path, '/models/fraud_model.pkl')
    
    print(f"Deployed model from {model_path}")
    
    # In production, you'd:
    # 1. Update model registry
    # 2. Restart API pods
    # 3. Send notification
    
    # Update model version in Prometheus
    # (You'd expose a metric for this)
    
    return model_path

def send_notification(**context):
    """Send notification about retraining"""
    new_accuracy = context['task_instance'].xcom_pull(
        task_ids='train_model',
        key='new_accuracy'
    )
    
    message = f"""
    🎉 Model Retraining Complete!
    
    New model deployed with accuracy: {new_accuracy:.2%}
    Timestamp: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}
    
    Check Grafana for updated metrics.
    """
    
    print(message)
    
    # In production, send to Slack:
    # requests.post(slack_webhook_url, json={'text': message})
    
    return message

# Define tasks
check_performance = PythonSensor(
    task_id='check_performance',
    python_callable=check_model_performance,
    mode='poke',
    poke_interval=300,  # Check every 5 minutes
    timeout=3600,       # Give up after 1 hour
    dag=dag
)

prepare_data = PythonOperator(
    task_id='prepare_data',
    python_callable=prepare_training_data,
    dag=dag
)

train_model = PythonOperator(
    task_id='train_model',
    python_callable=train_new_model,
    dag=dag
)

validate_model = PythonOperator(
    task_id='validate_model',
    python_callable=validate_new_model,
    dag=dag
)

deploy = PythonOperator(
    task_id='deploy_model',
    python_callable=deploy_model,
    dag=dag
)

notify = PythonOperator(
    task_id='send_notification',
    python_callable=send_notification,
    dag=dag
)

# Define task dependencies
check_performance >> prepare_data >> train_model >> validate_model >> deploy >> notify
```

**What this DAG does:**

1. **check_performance** - Queries Prometheus, decides if retraining needed
2. **prepare_data** - Pulls fresh training data (last 90 days)
3. **train_model** - Trains new model, logs to MLflow
4. **validate_model** - Ensures new model is better
5. **deploy** - Deploys new model to production
6. **notify** - Sends Slack notification

**The magic:** Fully automated. If accuracy drops, Airflow retrains without human intervention.

---

## **Part 5: The Complete Monitoring Workflow**

Let me show you how everything works together:

### **Normal Operation**
```
1. User makes API request
2. API makes prediction, records metrics
3. Prometheus scrapes metrics every 15s
4. Grafana displays real-time dashboards
5. Accuracy: 95% ✅
6. Latency: 80ms ✅
7. Everyone's happy
```

### **Performance Degradation**
```
Week 1: Accuracy 95%
Week 2: Accuracy 93%
Week 3: Accuracy 89% ⚠️
Week 4: Accuracy 86% 🚨

Prometheus Alert: "ModelAccuracyLow" fires
Airflow DAG: check_performance sensor detects issue
```

### **Automatic Retraining**
```
Day 1, 2:00 AM:
- Airflow checks Prometheus
- Accuracy: 86% < 90% threshold
- Decision: Retrain needed ✅

Day 1, 2:05 AM:
- Pull last 90 days of data
- 150,000 transactions loaded

Day 1, 2:15 AM:
- Train new Random Forest model
- New accuracy: 94%
- Log to MLflow

Day 1, 2:20 AM:
- Validate: 94% > 86% ✅
- Deploy new model
- Restart API

Day 1, 2:25 AM:
- Slack notification sent
- Grafana shows new model version
- Accuracy back to 94% ✅

python

# Accuracy (when ground truth available)
model_accuracy = Gauge('ml_model_accuracy', ...)

# Precision & Recall
model_precision = Gauge('ml_model_precision', ...)
model_recall = Gauge('ml_model_recall', ...)

# F1 Score
model_f1 = Gauge('ml_model_f1', ...)

# AUC-ROC
model_auc = Gauge('ml_model_auc', ...)

python

# Request rate
predictions_total = Counter('ml_predictions_total', ...)

# Latency percentiles
prediction_latency = Histogram('ml_prediction_latency_seconds', ...)

# Error rate
prediction_errors = Counter('ml_prediction_errors_total', ...)

# Throughput
predictions_per_second = Gauge('ml_predictions_per_second', ...)

python

# Feature statistics
feature_mean = Gauge('ml_feature_mean', ['feature_name'])
feature_std = Gauge('ml_feature_std', ['feature_name'])
feature_min = Gauge('ml_feature_min', ['feature_name'])
feature_max = Gauge('ml_feature_max', ['feature_name'])

# Data drift
prediction_distribution = Histogram('ml_prediction_distribution', ...)
feature_distribution = Histogram('ml_feature_distribution', ['feature'])

python

# Fraud detection rate
fraud_rate = Gauge('ml_fraud_rate', ...)

# False positive rate (when known)
false_positive_rate = Gauge('ml_false_positive_rate', ...)

# Revenue impact
revenue_protected = Counter('ml_revenue_protected_usd', ...)

python

# Only track predictions, no accuracy
predictions_total.inc()
# How do you know if model is still good?

python

# Track accuracy when ground truth arrives
@app.post("/feedback")
async def feedback(transaction_id, actual):
    predicted = get_prediction(transaction_id)
    is_correct = (predicted == actual)
    accuracy_metric.set(calculate_accuracy())

python

# Deploy and forget
# Hope model stays good forever
# Wait for users to complain

python

# Automated retraining in Airflow
# Triggered by performance metrics
# Validated before deployment
# Logged and versioned

python

# Only monitor model accuracy
# Ignore changing input distributions
# Miss early warning signs

python

# Track feature statistics
feature_stats.labels(
    feature='transaction_amount',
    stat='mean'
).set(features['transaction_amount'])

# Alert on drift
```

---

## **Production Checklist**

Before you deploy:

### **Metrics**
- ✅ Predictions total (by class)
- ✅ Latency percentiles (P50, P95, P99)
- ✅ Error rates
- ✅ Model accuracy (when available)
- ✅ Feature statistics
- ✅ Business metrics

### **Alerts**
- ✅ Accuracy drop alerts
- ✅ Latency spike alerts
- ✅ Error rate alerts
- ✅ Data drift alerts
- ✅ Alert routing (Slack, PagerDuty)

### **Dashboards**
- ✅ Real-time metrics dashboard
- ✅ Model performance trends
- ✅ Error investigation dashboard
- ✅ Business metrics dashboard

### **Retraining**
- ✅ Airflow DAG for retraining
- ✅ Performance threshold triggers
- ✅ Validation before deployment
- ✅ Rollback capability
- ✅ Notifications

### **Documentation**
- ✅ Runbook for common issues
- ✅ Alert response procedures
- ✅ Retraining process docs
- ✅ Dashboard explanations