Productionizing ML on Databricks with MLflow
MLflow solves experiment tracking and model versioning. It does not solve drift, retraining, or feature consistency. Here's the realistic path from notebook to production.
A model that scores well in a notebook and a model serving predictions in production are two different projects wearing the same accuracy metric. The gap between them is where most ML initiatives stall — not because the modeling was wrong, but because nobody planned for what happens after the notebook works. MLflow closes part of that gap. Here’s exactly which part, because the rest is still on you.
What MLflow actually solves
MLflow is not a magic production button. It’s three focused tools that fix three specific, common failures.
Experiment tracking solves the “which run was that?” problem. Every data scientist has a folder of notebooks named model_v3_final_REAL.ipynb and no reliable memory of which hyperparameters produced which metric. MLflow’s tracking API logs parameters, metrics, and artifacts for every run automatically, so “we got 0.91 AUC three weeks ago, what changed since” has an actual answer. On Databricks this is on by default for notebook training runs — no extra setup, which is most of why teams start here.
The model registry solves versioning and promotion. A registered model has a lineage: which run produced it, what stage it’s in (staging, production, archived), and who approved the transition. Without this, “roll back to last week’s model” means someone remembers which pickle file that was. With it, it’s one API call and an audit trail of every model that has ever served traffic.
Model packaging solves the “works on my cluster” problem. MLflow’s model format wraps the model with its dependencies, input schema, and inference logic into something that loads the same way whether you’re calling it from a batch job, a REST endpoint, or someone else’s notebook — portable instead of tied to the environment it was trained in.
All three of these are real, valuable, and solved cleanly by MLflow. None of them are “the model is in production.”
The gap MLflow doesn’t close
Here’s what a notebook that “works” is missing, in order of how often it bites teams:
Feature consistency between training and serving. The model was trained on features computed by a batch pipeline, joined and aggregated in ways that took real thought. In production, those same features often get recomputed by a different code path — an API service, a streaming job — written by someone who wasn’t in the room for the original feature engineering. Small discrepancies (a different null-handling rule, a stale join, a timezone off by one) silently degrade predictions with no error thrown. This is the single most common cause of “the model worked great in testing and then got worse in production,” and it’s an engineering problem, not a modeling one.
Monitoring for drift. A model’s accuracy on the day it ships is not its accuracy six months later. Input distributions shift — customer behavior changes, upstream systems change their data, seasonality happens — and a model trained on last year’s patterns quietly gets worse at this year’s problem. Without monitoring on prediction distributions and, where you can get it, actual outcomes, degradation looks like silence. Nobody gets paged because nothing crashed; the model just gets steadily less useful until someone notices the business metric moved.
Retraining triggers. Once you’re monitoring drift, you need a decision: at what threshold does this model get retrained, on what data window, and automatically or with a human sign-off? Teams that skip this either retrain on a fixed calendar regardless of need, or wait until someone escalates that the numbers look wrong — both are worse than a defined trigger tied to an actual drift metric.
Reliable, monitored serving. Wrapping a model in an endpoint and calling it done ignores latency under load, what happens when the endpoint is unhealthy, and whether a bad input crashes the request or degrades gracefully. Production serving needs the same operational rigor as any other production service — because it is one.
None of this is exotic. It’s the standard list of things that separate a working prototype from an operated system, applied to ML instead of a web app.
Unity Catalog’s role
If you’re already on Unity Catalog, the useful move is treating models and features as governed assets alongside your data, not a separate system living in someone’s workspace. Unity Catalog can hold your registered models with the same access controls, lineage, and audit logging as your tables — so “which model is serving predictions for the pricing endpoint, who approved it, and what data trained it” is answerable the same way you’d answer a data lineage question. Feature tables get the same treatment: governed, discoverable, and — critically — the same object whether a training job or a serving job reads them, which directly fixes the training/serving skew problem above. If you haven’t made that jump yet, we’ve written about what the Unity Catalog migration actually involves; for ML specifically, it’s the difference between features being a shared asset and duplicated logic waiting to drift apart.
How to scope a first production ML project
The projects that stall are the ones scoped as “productionize the model.” That’s not a scope, it’s a direction. A first production ML project that actually ships looks narrower:
- One model, one decision it feeds. Not a platform, not a suite of models — one prediction that changes one action someone takes. A churn score that triggers a retention workflow. A demand forecast that feeds a reorder point. Pick something where a wrong prediction is annoying, not catastrophic, for the first one.
- Batch before real-time, if you have the choice. A model that scores a table nightly and writes results back is dramatically simpler to operate than a low-latency endpoint. Get feature consistency and monitoring right in batch first; add real-time serving once you trust the pipeline.
- Define the retraining trigger before launch, not after the first complaint. Even a simple one — retrain monthly, or if a drift metric crosses a threshold — beats no plan, and it’s a smaller conversation to have before launch than after the model has been quietly wrong for two months.
- Instrument before you need it. Log predictions and, where feasible, eventual outcomes from day one. You cannot detect drift retroactively on data you didn’t keep.
Scoped this way, a first production ML project is weeks, not quarters, and it leaves you a working pattern — registry, feature consistency, monitoring, retraining trigger — to repeat for the second model, which is where the real payoff shows up.
This is the part of ML work that doesn’t show up in a demo, and it’s most of what determines whether a model earns its keep. Our data engineering and AI team builds this end-to-end on Databricks — tracking, registry, feature pipelines, monitoring — for teams moving their first models past the prototype stage. If you’ve got a model that works in a notebook and no clear path to production, get in touch and we’ll help you scope the first one right.
- Databricks
- AI