DEV Community

harshi606
harshi606

Posted on

I built an ML retraining pipeline that's allowed to say no

Here's a thing that bugged me while I was learning MLOps: almost every tutorial ends the same way. Train a model, print the accuracy, done. Cool, but... then what? What happens in three weeks when new data shows up? Does the model just get retrained and shoved into production because a cron job said so? What if the new one is actually worse?

Nobody seemed to want to answer that part. So I built something that does.

I come from a backend/software engineering background, not classical ML, so honestly the model itself was the least interesting part of this to me. I used a plain Random Forest. I'm not trying to win a Kaggle competition. What I actually wanted to figure out was: how do you build a system that can look at a freshly trained model and say "no, not good enough" — and mean it?

The actual problem

Say you've got a churn prediction model in production. A month goes by, you've got new customer data, and you retrain. Now you have a new model. Is it better? Maybe! Maybe not. Maybe it just overfit to some weird pocket of last month's data. If you don't have a real answer to "is it better," you shouldn't be letting it anywhere near production.

So I built a pipeline with an actual gate in it. A promotion gate. The new model has to prove itself before it's allowed to replace what's currently live.

How the gate works

Two conditions, both have to pass:

The new model has to clear a minimum bar on its own (I used F1 ≥ 0.70). Doesn't matter what's in production — if it's bad, it's rejected, full stop.
It has to beat whatever's currently in production, by a real margin. Not tie. Not "basically the same." I set the bar at 0.5% absolute F1 improvement.

That second one matters more than it sounds like it should. If you promote on literally any improvement, even 0.0001, you get this annoying situation where production flips back and forth based on random noise instead of anything meaningful. Requiring an actual margin stops that.

python
def promote_if_better(candidate_version, candidate_metrics, model_name):
candidate_score = candidate_metrics["f1"]

if candidate_score < MIN_ACCEPTABLE_F1:
    return {"promoted": False, "reason": "below the floor, not even close"}

prod_version, prod_metrics = get_current_production_metrics(model_name)

if prod_version is None:
    # nothing to compare against, this is the first one
    set_alias(model_name, "production", candidate_version)
    return {"promoted": True}

improvement = candidate_score - prod_metrics["f1"]

if improvement >= PROMOTION_MIN_IMPROVEMENT:
    set_alias(model_name, "previous-production", prod_version)  # keep it, just in case
    set_alias(model_name, "production", candidate_version)
    return {"promoted": True, "improvement": improvement}

return {"promoted": False, "improvement": improvement}
Enter fullscreen mode Exit fullscreen mode

That previous-production alias is doing quiet but important work. Before I let a new model take over, I tag the old one so I can find it again. That's what makes rollback a single function call instead of me digging through run history at 2am trying to remember which version was good.

I didn't just trust it, I tried to break it

Writing the logic is one thing. Believing it actually works is another. So before I let myself trust it, I ran it through a few scenarios by hand:

First model ever, nothing to compare to → should just promote. It did.
A model that ties with production → should get rejected. It did.
A model that's genuinely 3% better → should promote, and the old one should stick around for rollback → it did, and rollback actually worked.

Only after all three checked out did I go write proper integration tests for it — four of them, running against a real (but temporary) MLflow store so they're testing actual behavior, not just asserting that some constant is greater than zero.

python
def test_weak_candidate_is_rejected(mlflow_tracking):
v1 = register_and_promote(f1=0.80)
v2 = register_candidate(f1=0.801) # barely ties

decision = promote_if_better(v2, {"f1": 0.801})

assert decision["promoted"] is False
# and production should still be v1, unchanged
Enter fullscreen mode Exit fullscreen mode

Genuinely satisfying watching that test suite go green after knowing I'd already proven the logic by hand first.

Why the data has to drift

If you're not simulating some kind of change over time, there's no real reason to have a retraining pipeline at all — you'd just train once and stop. So I generated a synthetic churn dataset (sklearn's make_classification, nothing fancy) and then made a second "new batch" of data with a slight shift applied to it, like customer behavior nudging over time. That's the whole justification for the pipeline existing in the first place.

The part nobody puts in tutorials: it broke

Midway through wiring up Prefect to orchestrate all this, I got a error that looked like this:

Cool, cool, very helpful. Turns out the pipeline had actually finished — trained the model, registered it, made the promotion decision, all correctly — and then crashed while Prefect was cleaning up afterward. A version mismatch between Prefect and one of its dependencies, anyio. Prefect wanted anyio somewhere between 4.4.0 and 5.0.0, and whatever version pip grabbed inside that range had changed something internally that Prefect's code didn't expect.

Pinned it down to anyio==4.4.0 and it went away.

I'm including this because I think it's actually the most "real" part of the whole project. Nobody warns you that half of building these pipelines is just fighting with dependency versions that don't agree with each other. Reading a traceback carefully enough to realize "wait, this isn't my code, my code already finished" is its own skill, and it's one you only get by actually hitting the wall.

What I think this actually proves

Not "I can train a model." Anyone can train a model. What I think this shows instead:

I can build a system that's allowed to say no, not just one that says yes to everything
Every model version is traceable back to the exact run that produced it — no mystery models
Rollback isn't theoretical, I actually tested it
The whole thing runs as one command, not a pile of notebook cells I run in a specific order and hope I remember correctly
I can debug real infrastructure problems, not just tune hyperparameters

What's next

This pipeline is good at deciding which model deserves to be in production. It doesn't actually serve anything yet. So next up: wrapping the promoted model in a real API, with some actual monitoring on it, so I can watch it in the wild instead of just trusting it blindly. That's the other half of this — one system decides what's good enough, the other one keeps an eye on it once it's out there.

Code's here: https://github.com/harshi606/churn-retraining-pipeline

If you've built something similar, especially around the "how do I actually trust an automated promotion" problem, I'd genuinely like to hear how you approached it.

Top comments (0)