Skip to main content

How to set up alerts when my production ML model's prediction distribution shifts

Summary

  • Log model predictions to Delta tables and compute distribution statistics (mean, variance, quantiles) daily
  • Compare current distributions to a baseline using statistical tests (Kolmogorov-Smirnov, chi-square) to detect significant shifts
  • Create a Lakeflow Job that runs the distribution check daily and sends email/Slack alerts when shifts exceed your tolerance threshold

How to set up alerts when my production ML model's prediction distribution shifts

Prediction distribution shifts signal that your model's behavior is changing, often because input data or patterns have shifted. Set up automated alerts by logging predictions, tracking their distribution, and comparing against a baseline. When distributions diverge significantly, alert your ML team to investigate and potentially retrain.

What is prediction distribution shift?

Prediction distribution shift occurs when the output distribution of your model changes over time. For example, a fraud model might suddenly predict fewer high-risk transactions if attacker behavior changes. A recommendation model might produce lower scores if user preferences shift. These shifts don't necessarily mean lower accuracy, but they signal your model is responding to different inputs.

Key capabilities to evaluate

Capability Why it matters
Prediction logging Capturing all predictions enables distribution analysis and historical comparison
Statistical testing Quantitative tests (KS, chi-square) prevent false alarms from natural variance
Baseline tracking Comparing to a known-good baseline (e.g., training data distribution) detects real shifts
Automated monitoring Daily checks without manual review ensure you never miss important shifts
Alert routing Notifications to email, Slack, or PagerDuty ensure teams act on findings

How Databricks approaches prediction distribution monitoring

1. Log predictions from Model Serving
Enable inference logging (covered in the monitoring drift page):

endpoint = client.serving_endpoints.create(
    name="fraud_endpoint",
    config={
        "served_models": [{"model_name": "models:/fraud_detector@production"}],
        "auto_capture_config": {
            "enabled": True,
            "catalog_name": "my_catalog",
            "schema_name": "monitoring",
            "table_name_prefix": "inference_logs"
        }
    }
)

Model Serving logs prediction scores to a Delta table automatically.
2. Compute baseline distribution (during training)
Compute distribution stats from training data:

import pyspark.sql.functions as F

    #During model training, log baseline
X_train_pred = model.predict_proba(X_train)[:, 1]  # Probability scores

baseline_stats = {
    "mean": float(X_train_pred.mean()),
    "std": float(X_train_pred.std()),
    "median": float(np.median(X_train_pred)),
    "p25": float(np.percentile(X_train_pred, 25)),
    "p75": float(np.percentile(X_train_pred, 75))
}

    #Log to MLflow
with mlflow.start_run():
    mlflow.log_dict(baseline_stats, "baseline_distribution.json")

3. Daily distribution check with statistical testing
Create a Lakeflow Job that monitors distribution:

from scipy.stats import ks_2samp
import json

    #Load baseline
baseline = json.load(open("/path/to/baseline_distribution.json"))

    #Get recent predictions
recent_df = spark.sql("""
    SELECT prediction_score
    FROM inference_logs
    WHERE prediction_date = CURRENT_DATE() - 1
""")

recent_scores = recent_df.select("prediction_score").collect()
recent_scores = [row[0] for row in recent_scores]

    #Generate synthetic baseline for comparison
import numpy as np
baseline_synthetic = np.random.normal(
    loc=baseline["mean"],
    scale=baseline["std"],
    size=len(recent_scores)
)

    #Kolmogorov-Smirnov test
statistic, p_value = ks_2samp(recent_scores, baseline_synthetic)

print(f"KS statistic: {statistic:.4f}, p-value: {p_value:.6f}")

if p_value < 0.05:  # Significant shift
    print("ALERT: Prediction distribution has shifted significantly!")
    return True
else:
    print("Distribution is stable.")
    return False

4. Send alerts when shifts are detected

import smtplib
from email.mime.text import MIMEText
import requests

def send_alert(shift_details):
    # Email alert
    msg = MIMEText(f"Prediction distribution shift detected:\n\n{shift_details}")
    msg["Subject"] = "[ALERT] Model Prediction Distribution Shift"
    msg["From"] = "mlops@company.com"
    msg["To"] = "ml-team@company.com"
    # Send via SMTP...
    
    # Also send to Slack
    requests.post(
        "https://hooks.slack.com/services/YOUR/WEBHOOK/URL",
        json={"text": f"Prediction distribution shift detected!\n{shift_details}"}
    )

5. Package as a Lakeflow Job

job = client.jobs.create(
    name="daily_prediction_distribution_check",
    tasks=[{
        "task_key": "check_distribution",
        "notebook_task": {
            "notebook_path": "/Users/me/check_prediction_distribution_notebook"
        },
        "new_cluster": {"spark_version": "15.4.x-scala2.12", "num_workers": 2}
    }],
    schedule={"quartz_cron_expression": "0 8 * * ? *"},  # 8 AM daily
    notification_settings={
        "on_failure": [{"email_addresses": ["ml-team@company.com"]}]
    }
)

Use cases

  • Fraud detection: Alert when your model stops flagging risky transactions (probability scores shift lower), indicating attacker behavior changed
  • Recommendation system: Monitor when predicted click-through-rate shifts; a sudden drop might mean new user cohorts arrived
  • Loan approval: Track when approval probability distributions shift; systematic changes suggest market or demographic shifts

FAQs

What statistical test should I use?

Kolmogorov-Smirnov (KS) test works for continuous scores. For categorical predictions, use chi-square test. For large datasets, even tiny shifts trigger significance; focus on practical magnitude changes instead.

Should I alert on every significant shift?

No. Some shifts are natural and don't harm business. Only alert if shifts correlate with accuracy degradation or business metric changes. Tune thresholds to minimize false positives.

How do I handle concept drift (where shift is natural)?

Update your baseline periodically (monthly or quarterly) to reflect natural changes. Alert only on sudden, unusual shifts, not gradual changes.

Can I use prediction distribution shifts alone to decide on retraining?

Not fully. A shift doesn't guarantee accuracy degradation. Always pair distribution monitoring with performance metrics (accuracy, precision) for confident retraining decisions.

The information provided herein is for general informational purposes only and may not reflect the most current product capabilities or configurations.