Back to Articles

Deploying ML Models in Production: Lessons Learned

October 5, 2024
12 min read

The Gap Between Notebooks and Production

Every data scientist has built a model that works beautifully in a Jupyter notebook. The metrics look great, the stakeholders are excited, and you're ready to ship. Then reality hits.

Getting a model into production — and keeping it there — is an entirely different discipline. Google's famous paper "Hidden Technical Debt in Machine Learning Systems" illustrated this perfectly: the ML code is a tiny fraction of the actual system. The rest is plumbing, monitoring, and operational overhead.

Here are the lessons I've learned shipping ML models to production over the past few years.

Lesson 1: Your Model is Not the Product

The model is maybe 5% of the system. The rest is everything that keeps it running reliably:

  • Data pipelines — ingestion, validation, feature computation
  • Feature stores — consistent feature serving for training and inference
  • Serving infrastructure — APIs, batching, caching, load balancing
  • Monitoring — drift detection, latency tracking, error alerting
  • Retraining — automated pipelines triggered by performance degradation

Here's what this looks like in practice:

# What you think ML in production looks like:
prediction = model.predict(features)
 
# What it actually looks like:
def predict(user_id: str) -> dict:
    # 1. Fetch features with timeout and fallback
    try:
        features = feature_store.get_online_features(
            entity_key=user_id,
            feature_refs=["user:days_since_last_login",
                          "user:total_purchases_30d",
                          "user:avg_session_duration"],
            timeout_ms=100
        )
    except TimeoutError:
        logger.warning(f"Feature store timeout for {user_id}")
        return {"score": DEFAULT_SCORE, "source": "fallback"}
 
    # 2. Validate features
    if not validate_features(features):
        metrics.increment("prediction.invalid_features")
        return {"score": DEFAULT_SCORE, "source": "fallback"}
 
    # 3. Run inference
    score = model.predict_proba(features)[0][1]
 
    # 4. Log for monitoring
    prediction_logger.log({
        "user_id": user_id,
        "score": score,
        "features": features,
        "model_version": MODEL_VERSION,
        "timestamp": datetime.utcnow().isoformat()
    })
 
    # 5. Check for drift asynchronously
    drift_checker.submit(features)
 
    return {"score": round(score, 4), "source": "model"}

That's 40 lines for a single prediction. And this is the simplified version.

Lesson 2: Data Validation is Non-Negotiable

Bad data in production causes silent model failures. The model won't throw an error — it'll just return garbage predictions with full confidence.

Set up validation at every boundary:

from pydantic import BaseModel, validator
 
class PredictionRequest(BaseModel):
    user_id: str
    days_since_last_login: int
    total_purchases_30d: float
    avg_session_duration: float
 
    @validator("days_since_last_login")
    def validate_login_days(cls, v):
        if v < 0 or v > 3650:  # 10 years max
            raise ValueError(f"Invalid days_since_last_login: {v}")
        return v
 
    @validator("total_purchases_30d")
    def validate_purchases(cls, v):
        if v < 0:
            raise ValueError(f"Negative purchases: {v}")
        return v

I also recommend running Great Expectations or a similar tool on your training data pipeline:

import great_expectations as gx
 
context = gx.get_context()
validator = context.sources.pandas_default.read_csv("training_data.csv")
 
# Define expectations
validator.expect_column_values_to_not_be_null("user_id")
validator.expect_column_values_to_be_between("age", 13, 120)
validator.expect_column_mean_to_be_between("purchase_amount", 10, 500)
validator.expect_column_unique_value_count_to_be_between("country", 1, 250)
 
results = validator.validate()
if not results.success:
    raise DataQualityError(f"Training data validation failed: {results}")

Lesson 3: Monitor Everything (Especially Drift)

Model accuracy degrades silently. The world changes — customer behavior shifts, new products launch, seasonality kicks in — but your model is still optimizing for last quarter's patterns.

Set up monitoring across four dimensions:

Input Drift

Are the incoming features changing distribution compared to your training data?

from evidently.report import Report
from evidently.metric_preset import DataDriftPreset, TargetDriftPreset
 
# Run weekly drift check
drift_report = Report(metrics=[
    DataDriftPreset(),
    TargetDriftPreset()
])
 
drift_report.run(
    reference_data=training_df,
    current_data=last_week_predictions_df
)
 
# Extract results programmatically
drift_results = drift_report.as_dict()
drifted_features = [
    feature for feature, result
    in drift_results["metrics"][0]["result"]["drift_by_columns"].items()
    if result["drift_detected"]
]
 
if drifted_features:
    alert_slack(f"⚠️ Drift detected in: {', '.join(drifted_features)}")

Prediction Distribution

Are your model's outputs shifting? If your churn model suddenly predicts 80% churn for everyone, something is wrong — even if no individual prediction throws an error.

Latency

Track P50, P95, and P99 latency. ML models can slow down unexpectedly when feature computation gets expensive or when input sizes change.

Business Metrics

Ultimately, the model exists to move a business metric. Track the downstream impact:

  • If it's a recommendation model → click-through rate, conversion
  • If it's a churn model → retention rate, intervention success
  • If it's a fraud model → false positive rate, fraud losses

Lesson 4: Version Everything

Not just your model — version your entire ML pipeline:

# model_manifest.yaml
model:
  name: churn_predictor
  version: 2.3.1
  framework: xgboost
  created_at: "2024-09-15T10:30:00Z"
 
training:
  data_version: "s3://ml-data/churn/v12/train.parquet"
  data_hash: "sha256:a1b2c3d4..."
  rows: 245000
  features:
    - days_since_last_login
    - total_purchases_30d
    - support_tickets_open
    - avg_session_duration
    - account_age_days
 
hyperparameters:
  max_depth: 6
  learning_rate: 0.1
  n_estimators: 200
  subsample: 0.8
  colsample_bytree: 0.8
 
evaluation:
  auc_roc: 0.847
  precision_at_10pct: 0.72
  recall_at_10pct: 0.58
  test_set: "s3://ml-data/churn/v12/test.parquet"
 
serving:
  threshold: 0.65
  fallback_score: 0.5
  max_latency_ms: 100
  min_feature_coverage: 0.8

This manifest travels with the model artifact. When something goes wrong in production, you can trace back to the exact data, code, and hyperparameters that produced it.

Tools that help: MLflow for experiment tracking, DVC for data versioning, Weights & Biases for experiment visualization.

Lesson 5: Start Simple, Then Iterate

Don't jump to Kubernetes + Kubeflow + feature stores + A/B testing platform on day one. I've seen teams spend 6 months building infrastructure before serving a single prediction.

Here's the progression I recommend:

V1: Just Ship It (Week 1-2)

  • Flask/FastAPI endpoint wrapping your model
  • Cron job for weekly retraining
  • Basic logging to a file or database
  • Manual deployment
# v1: This is fine for starting out
from fastapi import FastAPI
import joblib
 
app = FastAPI()
model = joblib.load("model.pkl")
 
@app.post("/predict")
def predict(features: dict):
    score = model.predict_proba([list(features.values())])[0][1]
    return {"score": round(score, 4)}

V2: Add Observability (Month 1-2)

  • Structured logging with prediction IDs
  • Basic drift monitoring (weekly Evidently reports)
  • Alerting on error rates and latency
  • Automated retraining pipeline (Airflow or similar)

V3: Scale and Harden (Month 3-6)

  • Feature store for consistent feature computation
  • Model registry with staged deployments (staging → production)
  • Shadow mode for new models (run in parallel, compare outputs)
  • Canary deployments (10% → 50% → 100%)

V4: Optimize (Month 6+)

  • A/B testing framework for model comparisons
  • Online learning or more frequent retraining
  • Advanced drift detection with automated retraining triggers
  • Cost optimization (batching, caching, model compression)

Lesson 6: Have a Fallback Strategy

Your model will fail in production. Network timeouts, feature store outages, corrupted model artifacts, OOM errors — the question isn't if, but when.

Design your system to fail gracefully:

class PredictionService:
    def __init__(self):
        self.model = load_model()
        self.fallback_rules = load_rules()  # Simple heuristics
        self.cache = PredictionCache(ttl=3600)
 
    def predict(self, user_id: str) -> PredictionResult:
        # Strategy 1: Try the ML model
        try:
            return self._ml_predict(user_id)
        except ModelError as e:
            logger.error(f"Model failed: {e}")
            metrics.increment("prediction.model_failure")
 
        # Strategy 2: Try cached prediction
        cached = self.cache.get(user_id)
        if cached:
            metrics.increment("prediction.cache_hit")
            return cached
 
        # Strategy 3: Rule-based fallback
        metrics.increment("prediction.rule_fallback")
        return self.fallback_rules.predict(user_id)

The key principle: the user should never see a failure. They might get a slightly less accurate prediction, but the feature should always work.

Lesson 7: Automate Retraining (But Don't Overtrain)

Retraining should be triggered by signals, not just schedules:

  • Drift detected → retrain on recent data
  • Performance dropped below threshold → retrain and investigate
  • New labeled data available → retrain with expanded dataset
  • Scheduled → weekly or monthly as a baseline

But be careful with automatic retraining. I've seen models degrade because they were retrained on corrupted data without validation. Always run your test suite on the new model before promoting it:

def retrain_pipeline():
    # 1. Pull latest data
    train_df, test_df = pull_training_data()
 
    # 2. Validate data quality
    assert validate_data(train_df), "Training data failed validation"
 
    # 3. Train new model
    new_model = train(train_df)
 
    # 4. Evaluate against test set
    metrics = evaluate(new_model, test_df)
 
    # 5. Compare against current production model
    prod_metrics = get_production_metrics()
 
    if metrics["auc_roc"] < prod_metrics["auc_roc"] - 0.02:
        alert("New model performs worse than production. Skipping deployment.")
        return
 
    # 6. Deploy to staging first
    deploy_to_staging(new_model)
    run_shadow_comparison(hours=24)
 
    # 7. Promote to production if shadow mode looks good
    if shadow_comparison_passed():
        promote_to_production(new_model)
        alert(f"New model v{new_model.version} deployed. AUC: {metrics['auc_roc']:.3f}")

Key Takeaways

  1. The model is 5% of the system — invest more in infrastructure than model tuning
  2. Validate data at every boundary — bad data causes silent failures
  3. Monitor four dimensions — input drift, prediction drift, latency, business metrics
  4. Version everything — data, code, config, hyperparameters, and model artifacts
  5. Start simple — Flask API → observability → feature store → A/B testing
  6. Always have a fallback — ML model → cache → rules → default value
  7. Automate retraining carefully — validate before deploying, compare against production

The best ML system is one that fails gracefully, recovers automatically, and tells you when it's getting worse. Build for reliability first, accuracy second. A model that's 2% less accurate but never goes down is infinitely more valuable than a model that's state-of-the-art but breaks at 3 AM on a Friday.

Back to All Articles