Skip to main content
MLOps · lesson 4 of 8
MLOps intermediate Lesson 4 of 8

Also part of: Machine Learning

Model Serving and Deployment

Serve a model with FastAPI and pydantic, load it once from an MLflow alias, batch-score Parquet in chunks, and gate promotion and rollback on a holdout.

By SkillByExample · Updated

Real-World Scenario

A churn model starts as a notebook and a pickle file emailed to the backend team. Three months later nobody knows which file production loaded, the service crashes on a customer with negative tenure from a data glitch, and the nightly batch job scores with a different model from the API. Serving is the discipline that fixes this. Load one registered version by name, validate every input before it reaches the model, use the same version for online and batch scoring, and change versions only through a gate that can say no.

Every example here runs locally with FastAPI, pydantic, uvicorn and MLflow 3’s default local store in the current directory, with no server and no cloud account: pip install fastapi uvicorn mlflow scikit-learn pandas pyarrow. The experiment tracking lesson covers the MLflow registry and aliases in more depth.

A Model in the Registry

The data and feature code is shared by training, serving and batch jobs, so it lives in one file, data.py. A feature computed differently in training and in serving is a common production bug, and one shared function prevents it.

import numpy as np
import pandas as pd

CONTRACTS = ["month-to-month", "one-year", "two-year"]

def customers(n, seed=0):
    rng = np.random.default_rng(seed)
    df = pd.DataFrame({
        "tenure_months": rng.integers(0, 72, n),
        "monthly_charges": rng.uniform(20, 120, n).round(2),
        "contract": rng.choice(CONTRACTS, n, p=[0.5, 0.3, 0.2]),
        "support_calls_30d": rng.poisson(1.2, n),
    })
    risk = (0.03 * df.monthly_charges - 0.06 * df.tenure_months + 3.0 * (df.support_calls_30d >= 3)
            - 1.5 * df.contract.map({"month-to-month": 0, "one-year": 1, "two-year": 2}))
    df["churned"] = (risk + rng.normal(0, 3, n) > 0).astype(int)
    return df

def features(df):
    X = df[["tenure_months", "monthly_charges", "support_calls_30d"]].copy()
    X["contract"] = df.contract.map({c: i for i, c in enumerate(CONTRACTS)})
    return X.astype(float)

train.py fits a logistic regression, registers it as churn-service, and points the alias champion at it:

import mlflow
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import roc_auc_score
from data import customers, features

train, test = customers(5_000, seed=1), customers(1_000, seed=2)
model = LogisticRegression(max_iter=1_000).fit(features(train), train.churned)
auc = roc_auc_score(test.churned, model.predict_proba(features(test))[:, 1])

mlflow.set_experiment("churn-serving")
with mlflow.start_run(run_name="logreg-v1"):
    mlflow.log_metric("auc", auc)
    info = mlflow.sklearn.log_model(model, name="model", registered_model_name="churn-service")
mlflow.MlflowClient().set_registered_model_alias("churn-service", "champion",
                                                 info.registered_model_version)
print(f"test AUC {auc:.4f}; churn-service v{info.registered_model_version} is @champion")
$ python train.py
2026/09/29 18:39:20 INFO mlflow.store.db.utils: Creating initial MLflow database tables...
2026/09/29 18:39:20 INFO mlflow.store.db.utils: Updating database tables
2026/09/29 18:39:21 INFO mlflow.tracking.fluent: Experiment with name 'churn-serving' does not exist. Creating a new experiment.
Successfully registered model 'churn-service'.
Created version '1' of model 'churn-service'.
test AUC 0.7758; churn-service v1 is @champion

MLflow created mlflow.db in the current directory and registered version 1. From here on, nothing refers to a file path or a version number: the service asks for models:/churn-service@champion, and whoever controls the alias controls what production runs.

The Prediction API

The request and response are pydantic models, saved as schemas.py. Every constraint here is checked before your code runs.

from typing import Literal
from pydantic import BaseModel, Field

class Customer(BaseModel):
    customer_id: str
    tenure_months: int = Field(ge=0, le=600)
    monthly_charges: float = Field(ge=0, le=10_000)
    contract: Literal["month-to-month", "one-year", "two-year"]
    support_calls_30d: int = Field(ge=0)

class Prediction(BaseModel):
    customer_id: str
    churn_probability: float
    risk: Literal["LOW", "MEDIUM", "HIGH"]
    model_version: str

serve.py resolves the champion alias to a version number once, loads that exact version in FastAPI’s lifespan startup hook, and reuses it for every request. Resolving first means model_version in every response names the model actually loaded, even if someone moves the alias while the server is starting:

from contextlib import asynccontextmanager
import mlflow
import pandas as pd
from fastapi import FastAPI
from data import features
from schemas import Customer, Prediction

state = {}

@asynccontextmanager
async def lifespan(app):
    version = mlflow.MlflowClient().get_model_version_by_alias("churn-service", "champion").version
    state["model"] = mlflow.sklearn.load_model(f"models:/churn-service/{version}")  # once, at startup
    state["version"] = str(version)                # the exact version loaded, even if the alias moves
    yield

app = FastAPI(title="churn-service", lifespan=lifespan)

@app.post("/predict", response_model=Prediction)
def predict(customer: Customer):
    X = features(pd.DataFrame([customer.model_dump()]))
    p = float(state["model"].predict_proba(X)[0, 1])
    risk = "HIGH" if p >= 0.7 else "MEDIUM" if p >= 0.4 else "LOW"
    return Prediction(customer_id=customer.customer_id, churn_probability=round(p, 4),
                      risk=risk, model_version=state["version"])

@app.get("/health")
def health():
    return {"status": "ok", "model_version": state.get("version")}

Start it with uvicorn serve:app --port 8000 in the same directory, then call it from a second terminal: a high-risk customer, a low-risk one, and a request with two impossible values.

$ curl -s localhost:8000/health
{"status":"ok","model_version":"1"}
$ curl -s localhost:8000/predict -H 'content-type: application/json' -d '{"customer_id": "c-1", "tenure_months": 6, "monthly_charges": 85, "contract": "month-to-month", "support_calls_30d": 3}'
{"customer_id":"c-1","churn_probability":0.8873,"risk":"HIGH","model_version":"1"}
$ curl -s localhost:8000/predict -H 'content-type: application/json' -d '{"customer_id": "c-2", "tenure_months": 60, "monthly_charges": 35, "contract": "two-year", "support_calls_30d": 0}'
{"customer_id":"c-2","churn_probability":0.0348,"risk":"LOW","model_version":"1"}
$ curl -s -w '\nHTTP %{http_code}\n' localhost:8000/predict -H 'content-type: application/json' -d '{"customer_id": "c-3", "tenure_months": -3, "monthly_charges": 99.5, "contract": "weekly", "support_calls_30d": 4}'
{"detail":[{"type":"greater_than_equal","loc":["body","tenure_months"],"msg":"Input should be greater than or equal to 0","input":-3,"ctx":{"ge":0}},{"type":"literal_error","loc":["body","contract"],"msg":"Input should be 'month-to-month', 'one-year' or 'two-year'","input":"weekly","ctx":{"expected":"'month-to-month', 'one-year' or 'two-year'"}}]}
HTTP 422

$ cat server.log
INFO:     Started server process [15659]
INFO:     Waiting for application startup.
INFO:     Application startup complete.
INFO:     Uvicorn running on http://127.0.0.1:8000 (Press CTRL+C to quit)
INFO:     127.0.0.1:33220 - "GET /health HTTP/1.1" 200 OK
INFO:     127.0.0.1:33232 - "GET /health HTTP/1.1" 200 OK
INFO:     127.0.0.1:33242 - "POST /predict HTTP/1.1" 200 OK
INFO:     127.0.0.1:33250 - "POST /predict HTTP/1.1" 200 OK
INFO:     127.0.0.1:33252 - "POST /predict HTTP/1.1" 422 Unprocessable Entity
INFO:     Shutting down
INFO:     Waiting for application shutdown.
INFO:     Application shutdown complete.
INFO:     Finished server process [15659]

A new customer on a monthly contract with three support calls scores 0.8873, HIGH. A five-year customer on a two-year contract scores 0.0348, LOW. Both responses carry model_version, so any logged prediction can be traced back to the exact registered model. The third request never reached the model. pydantic rejected it with HTTP 422 and listed both problems at once, negative tenure and a contract type that doesn’t exist, which is what the caller needs to fix their request. Without validation, -3 alone would have been scored and returned as a confident but meaningless number, and "weekly" alone would have become a missing value in features and crashed the handler with a 500 instead of a useful 400-class error. The log was printed after the server stopped, so it ends with the shutdown lines, and its extra /health call is the script polling until the server was up.

The Failure: Loading the Model Per Request

It’s tempting to write model = mlflow.sklearn.load_model(...) inside predict, so every request “gets the latest model”. Measure what that costs: the first load in a process, later loads once everything is warm, and one prediction.

import time
import mlflow
from data import customers, features

X = features(customers(1, seed=9))

def timed_load():
    start = time.perf_counter()
    model = mlflow.sklearn.load_model("models:/churn-service@champion")
    return model, (time.perf_counter() - start) * 1000

model, cold_ms = timed_load()                      # first load in the process
warm_ms = min(timed_load()[1] for _ in range(5))   # what a per-request load would pay

start = time.perf_counter()
for _ in range(100):
    model.predict_proba(X)
predict_ms = (time.perf_counter() - start) * 1000 / 100

print(f"first load:        {cold_ms:8.1f} ms")
print(f"later loads:       {warm_ms:8.1f} ms  (best of 5)")
print(f"one prediction:    {predict_ms:8.2f} ms")
print(f"load per request would make each call ~{warm_ms / predict_ms:,.0f}x slower")
$ python latency.py
first load:          2543.9 ms
later loads:           20.0 ms  (best of 5)
one prediction:        1.46 ms
load per request would make each call ~14x slower

The first load took 2.5 seconds, but most of that is one-off work: importing MLflow’s scikit-learn support and opening the tracking database. Measuring only that would exaggerate the problem. Once warm, a load still took 20 ms, resolving the alias, reading the files and deserializing them, against 1.46 ms for the prediction itself. So loading in the handler makes every request about 14 times slower, and under load the server spends most of its time re-reading the same files; with a larger model the load grows and the ratio gets worse. It also means a request can start on one model version and the next on another, mid-traffic. Load at startup, as lifespan does. To pick up a new champion, restart or reload the servers, which makes the change a deliberate, observable event.

Batch Scoring a Parquet File

Batch inference scores a whole table, usually too large to hold in memory at once. First, write 200,000 customers to score, and try the obvious chunked read.

import pandas as pd
from data import customers

df = customers(200_000, seed=5).drop(columns="churned")
df.insert(0, "customer_id", [f"c-{i}" for i in range(len(df))])
df.to_parquet("to_score.parquet")
print(f"wrote {len(df):,} customers to to_score.parquet")

try:
    for chunk in pd.read_parquet("to_score.parquet", chunksize=50_000):   # like read_csv?
        print(len(chunk))
except TypeError as e:
    print("TypeError:", e)
$ python batch_naive.py
wrote 200,000 customers to to_score.parquet
TypeError: read_table() got an unexpected keyword argument 'chunksize'

pd.read_csv accepts chunksize; pd.read_parquet doesn’t. It passes the argument on to pyarrow’s read_table, which rejects it. The error names read_table, not read_parquet, which makes it harder to trace. Read in batches with pyarrow directly:

import time
import mlflow
import pandas as pd
import pyarrow.parquet as pq
from data import features

model = mlflow.sklearn.load_model("models:/churn-service@champion")
start, parts = time.perf_counter(), []
for i, batch in enumerate(pq.ParquetFile("to_score.parquet").iter_batches(batch_size=50_000)):
    chunk = batch.to_pandas()
    p = model.predict_proba(features(chunk))[:, 1]
    parts.append(pd.DataFrame({"customer_id": chunk.customer_id, "churn_probability": p.round(4),
                               "risk": pd.cut(p, [0, 0.4, 0.7, 1.01], labels=["LOW", "MEDIUM", "HIGH"],
                                              right=False)}))
    print(f"batch {i + 1}: {len(chunk):,} rows")

scores = pd.concat(parts, ignore_index=True)
scores.to_parquet("scores.parquet")
print(f"scored {len(scores):,} rows in {time.perf_counter() - start:.2f}s")
print(scores.risk.value_counts().sort_index().to_string())
$ python batch.py
batch 1: 50,000 rows
batch 2: 50,000 rows
batch 3: 50,000 rows
batch 4: 50,000 rows
scored 200,000 rows in 0.46s
risk
LOW       97620
MEDIUM    73679
HIGH      28701

iter_batches reads 50,000 rows at a time, so memory stays bounded however large the file is. The job loads the same @champion alias as the API and calls the same features function, so a customer gets the same score from the nightly batch as from a live request. Scoring 200,000 rows took under half a second, a fraction of the model’s first load, which is why batch jobs load once and score everything. The risk bands use the same 0.4 and 0.7 cut-offs as the API, and right=False makes each band include its lower edge, so a probability of exactly 0.7 is HIGH here just as p >= 0.7 makes it HIGH in the API. The default, right-closed bins would put it in MEDIUM, and the two paths would disagree at the boundary. In a larger system those thresholds belong in one shared place, like features.

Promoting a Challenger Through a Gate

A new model version should replace the champion only after it proves itself on data neither model trained on. challenger.py trains a random forest and registers it under the alias challenger. With --buggy, it simulates a pipeline bug that shuffled the labels.

import sys
import mlflow
from sklearn.ensemble import RandomForestClassifier
from data import customers, features

train = customers(5_000, seed=1)
if sys.argv[1:] == ["--buggy"]:                     # a pipeline bug: labels shuffled
    train["churned"] = train.churned.sample(frac=1, random_state=0).to_numpy()
model = RandomForestClassifier(n_estimators=100, max_depth=6, random_state=0)
model.fit(features(train), train.churned)

mlflow.set_experiment("churn-serving")
with mlflow.start_run(run_name="forest" + ("-buggy" if sys.argv[1:] else "")):
    info = mlflow.sklearn.log_model(model, name="model", registered_model_name="churn-service",
                                    skops_trusted_types=["sklearn.tree._tree.Tree"])
mlflow.MlflowClient().set_registered_model_alias("churn-service", "challenger",
                                                 info.registered_model_version)
print(f"churn-service v{info.registered_model_version} is @challenger")

promote.py scores both aliases on the same holdout set and moves champion only if the challenger is at least as good:

import mlflow
from sklearn.metrics import roc_auc_score
from data import customers, features

client, holdout = mlflow.MlflowClient(), customers(2_000, seed=3)

def auc(alias):
    model = mlflow.sklearn.load_model(f"models:/churn-service@{alias}")
    return roc_auc_score(holdout.churned, model.predict_proba(features(holdout))[:, 1])

champion, challenger = auc("champion"), auc("challenger")
version = client.get_model_version_by_alias("churn-service", "challenger").version
print(f"champion AUC {champion:.4f} | challenger v{version} AUC {challenger:.4f}")
if challenger >= champion - 0.005:                 # at least as good, within noise
    client.set_registered_model_alias("churn-service", "champion", version)
    print(f"promoted: @champion -> v{version}")
else:
    print(f"rejected: @champion stays on v{client.get_model_version_by_alias('churn-service', 'champion').version}")

Run the buggy challenger through the gate first, then a correct one.

$ python challenger.py --buggy
Registered model 'churn-service' already exists. Creating a new version of this model...
Created version '2' of model 'churn-service'.
churn-service v2 is @challenger

$ python promote.py
champion AUC 0.7473 | challenger v2 AUC 0.4471
rejected: @champion stays on v1

$ python challenger.py
Registered model 'churn-service' already exists. Creating a new version of this model...
Created version '3' of model 'churn-service'.
churn-service v3 is @challenger

$ python promote.py
champion AUC 0.7473 | challenger v3 AUC 0.7522
promoted: @champion -> v3

The buggy version registered without complaint. The registry stores whatever it is given, and training “succeeded”. Only the gate caught it: AUC 0.4471 is worse than a coin flip, and @champion stayed on v1. The correct forest scores 0.7522 on the holdout against the champion’s 0.7473, so it’s promoted. The next server start and the next batch run pick up v3 with no code change. The small 0.005 tolerance lets an equivalent model through, so the gate doesn’t block a retrain on fresh data over noise, but it never lets through a clear regression. In CI, run this script after training and fail the pipeline on “rejected”.

Practice

  1. Restart the server after the promotion and call /health. Then send a request that leaves out contract and support_calls_30d. What version is served, and what does the error say?

    Solution
    $ curl -s localhost:8000/health
    {"status":"ok","model_version":"3"}
    $ curl -s localhost:8000/predict -H 'content-type: application/json' -d '{"customer_id": "c-1", "tenure_months": 6, "monthly_charges": 85}'
    {"detail":[{"type":"missing","loc":["body","contract"],"msg":"Field required","input":{"customer_id":"c-1","tenure_months":6,"monthly_charges":85}},{"type":"missing","loc":["body","support_calls_30d"],"msg":"Field required","input":{"customer_id":"c-1","tenure_months":6,"monthly_charges":85}}]}

    The restarted server resolved @champion at startup and now serves v3, with no change to serve.py. A server started before the promotion would still report v1 until it restarts, so /health is how a deploy script checks that every instance picked up the new version. The missing fields fail with missing / Field required, both listed in one response. Fields without a default are required, and pydantic doesn’t guess values for them.

  2. The forest turns out to misbehave on a customer segment the holdout didn’t cover. Roll @champion back to v1 and list every version with its aliases and the run that produced it.

    Solution
    import mlflow
    
    client = mlflow.MlflowClient()
    before = client.get_model_version_by_alias("churn-service", "champion").version
    client.set_registered_model_alias("churn-service", "champion", 1)
    after = client.get_model_version_by_alias("churn-service", "champion").version
    print(f"@champion: v{before} -> v{after}")
    for n in (1, 2, 3):
        v = client.get_model_version("churn-service", n)
        print(f"v{n}  aliases={v.aliases}  run={client.get_run(v.run_id).info.run_name}")
    $ python rollback.py
    @champion: v3 -> v1
    v1  aliases=['champion']  run=logreg-v1
    v2  aliases=[]  run=forest-buggy
    v3  aliases=['challenger']  run=forest

    Rollback is the same one call as promotion, pointed at an older version. That only works because v1 was never deleted. Every version keeps its link to the run that produced it, so the buggy v2 stays on record for the post-mortem. v3 still carries challenger, because an alias moves only when someone moves it. Restart the servers to finish the rollback.

  3. AUC only measures ranking. Extend the gate so the challenger must also not raise the share of HIGH-risk predictions on the holdout by more than 5 percentage points, then run it against the current champion and challenger.

    Solution
    import mlflow
    from sklearn.metrics import roc_auc_score
    from data import customers, features
    
    holdout = customers(2_000, seed=3)
    
    def check(alias):
        model = mlflow.sklearn.load_model(f"models:/churn-service@{alias}")
        p = model.predict_proba(features(holdout))[:, 1]
        return roc_auc_score(holdout.churned, p), (p >= 0.7).mean()
    
    (auc_a, high_a), (auc_b, high_b) = check("champion"), check("challenger")
    print(f"champion   AUC {auc_a:.4f}  HIGH share {high_a:.1%}")
    print(f"challenger AUC {auc_b:.4f}  HIGH share {high_b:.1%}")
    ok = auc_b >= auc_a - 0.005 and high_b - high_a <= 0.05
    print("promote" if ok else "reject", f"(HIGH share moved {high_b - high_a:+.1%})")
    $ python prac_gate.py
    champion   AUC 0.7473  HIGH share 14.1%
    challenger AUC 0.7522  HIGH share 14.0%
    promote (HIGH share moved -0.1%)

    After the rollback, champion is v1 and challenger is v3 again. The forest passes both checks: it ranks slightly better and sends almost exactly the same share of customers, 14.0% against 14.1%, to the HIGH queue. A model can match the champion’s AUC while being calibrated very differently, which would double the retention team’s workload overnight without anyone having approved it. Real gates combine a quality metric with the things downstream consumers depend on: the share of each risk band, latency, and the input schema.

Frequently Asked Questions

What is the difference between online serving and batch inference?
Online serving answers one request at a time over HTTP, within milliseconds, for features such as fraud checks at checkout. Batch inference scores a whole table on a schedule and writes the results somewhere for other systems to read, such as a nightly churn score for every customer. Both should load the same registered model version so their predictions agree.
Why should a model be loaded at startup instead of in the request handler?
Loading a model is much slower than using it. In a measured example, reloading a small scikit-learn model from an MLflow registry took about 20 ms once warm, and the first load 2.5 seconds, while one prediction took about 1.5 ms. Loading inside the handler would make every request roughly 14 times slower, so load once in the application's startup hook and keep the model in memory.
How does FastAPI validate prediction requests?
Declare the request body as a pydantic model with types and constraints, such as Field(ge=0) for a count or a Literal for an allowed set of strings. FastAPI rejects a request that breaks them with HTTP 422 and a JSON list naming every failing field, before your prediction code runs. That keeps impossible values like negative tenure away from the model.
How do I read a large Parquet file in chunks with pandas?
You can't with pandas.read_parquet: it has no chunksize argument, and passing one fails with a TypeError from the underlying engine. Use pyarrow instead: pyarrow.parquet.ParquetFile(path).iter_batches(batch_size=50_000) yields record batches, and batch.to_pandas() turns each into a DataFrame to score and append.
How do I roll back a bad model in production?
If serving code loads the model by an alias such as models:/churn-service@champion, rolling back is one call that points the alias at the previous version, followed by a restart or reload of the servers. No code or image changes. Keep the old versions registered so there is always something to roll back to.