Skip to main content

Module 3 — Loading the model at startup

The placeholder score from modules 1 and 2 has done its job. This module wires the real churn model produced by course 20 into the service and, in doing so, fixes the biggest performance and reliability trap of naive ML APIs: loading the model somewhere that runs on every request. Doing it once, at process start, is what makes latency predictable and outages loud rather than silent.

Why per-request loading is a bug, not a style​

Consider the naive version that appears in almost every first draft:

# NEVER DO THIS
from fastapi import FastAPI
import mlflow.pyfunc

app = FastAPI()


@app.post("/predict")
def predict(payload: dict):
    model = mlflow.pyfunc.load_model("models:/churn-classifier@production")  # per request
    return {"probability": float(model.predict([payload])[0])}

Every request pays the cost of downloading (or reading from disk), deserializing, and initializing the model. For a scikit-learn gradient boosting the model, that is between 200 ms and 3 s per request; for a 400 MB transformer it is 5 to 30 s. Users see a p50 that is dominated by the load, not by the inference. Worse, if the registry is briefly unavailable — a network blip, an alias being moved — every request in that window fails, even though the process has a perfectly usable model in memory somewhere else.

The fix is architectural: load the model once, at process start, and keep it in memory for the life of the process. FastAPI has a first-class hook for exactly that.

lifespan: the modern way to load at startup​

Since Starlette 0.13 and FastAPI 0.93, the recommended way to run code at startup and shutdown is the lifespan context manager. It replaces the older @app.on_event("startup") decorator, which is now deprecated:

# app/main.py
import os
from contextlib import asynccontextmanager
from fastapi import FastAPI
import mlflow.pyfunc

from app.schemas import Subscriber, Prediction

MODEL_URI = os.environ.get("MODEL_URI", "models:/churn-classifier@production")
MODEL_STATE: dict = {}


@asynccontextmanager
async def lifespan(app: FastAPI):
    # Startup: load once. Any exception here aborts the process, which is
    # exactly what we want — the container's orchestrator will restart it
    # and mark the pod NotReady until the model is available.
    MODEL_STATE["model"] = mlflow.pyfunc.load_model(MODEL_URI)
    MODEL_STATE["uri"] = MODEL_URI
    MODEL_STATE["version"] = _resolve_version(MODEL_URI)
    yield
    # Shutdown: nothing to clean up for a pyfunc model, but the block is
    # where you would close DB pools, flush metrics, cancel background tasks.
    MODEL_STATE.clear()


app = FastAPI(title="Churn scoring API", version="0.3.0", lifespan=lifespan)


@app.post("/predict", response_model=Prediction, tags=["prediction"])
def predict(subscriber: Subscriber) -> Prediction:
    model = MODEL_STATE["model"]
    proba = float(model.predict([subscriber.model_dump()])[0])
    bucket = "low" if proba < 0.33 else "medium" if proba < 0.66 else "high"
    return Prediction(
        subscriber_id=subscriber.subscriber_id,
        churn_probability=round(proba, 4),
        risk_bucket=bucket,
        model_version=MODEL_STATE["version"],
    )

Everything before the yield runs on startup; everything after runs on shutdown. The MODEL_STATE dictionary is a plain module-level object, so every worker keeps its own copy — which matters for the memory sizing discussed in module 10.

_resolve_version is a small helper that turns an alias-based URI into the concrete version number, so the run that answered a specific request can be traced back to a specific model run in MLflow:

# app/versioning.py
from mlflow.tracking import MlflowClient


def _resolve_version(model_uri: str) -> str:
    # "models:/churn-classifier@production" -> concrete version like "17"
    if not model_uri.startswith("models:/"):
        return model_uri  # e.g. a file:// URI, treated as its own identity
    _, rest = model_uri.split("models:/", 1)
    name, at, alias = rest.partition("@")
    if not at:
        return rest  # already "name/version"
    mv = MlflowClient().get_model_version_by_alias(name=name, alias=alias)
    return f"{name}/{mv.version}"

Two lessons live inside this helper. First, the URI a caller sees at deploy time is not the URI a request executed against: @production moves, versions do not. Log the resolved version, not the alias. Second, the resolution runs at startup, not per request. If MLflow disappears for a minute during a deploy, in-flight requests still work; only new pods coming up fail loudly.

Fail fast on startup, not on the first request​

The point of loading at startup is not just performance. It is the fail-fast property. If the model is missing, the container will not enter the "ready" state, and the orchestrator will not send traffic to it. Compare with the naive version, where the first request after a bad deploy returns a 500 and everything looks fine until the user hits the wrong pod.

Concretely, in the code above:

  • mlflow.pyfunc.load_model raises on a missing artifact. The exception bubbles out of lifespan and the process exits with a non-zero code.
  • Kubernetes, ECS, or a plain docker run with --restart=on-failure restarts the container. The readiness probe (module 8) keeps it out of the load balancer until it comes up clean.
  • CI can boot the image with a dummy MODEL_URI and assert that the process fails within 5 seconds, catching bad artifacts before they ship.

None of these properties exist when the model is loaded on the first request.

Exposing the model version​

An ML service that cannot answer "which model did you just use?" is unauditable. Add a small, unauthenticated version endpoint:

# app/main.py (continued)
@app.get("/model", tags=["meta"])
def model_meta() -> dict:
    return {
        "uri": MODEL_STATE["uri"],
        "version": MODEL_STATE["version"],
        "service_version": app.version,
    }

This route pairs with a request-scoped log line (module 8) that includes the same version, so a support ticket that says "the customer got 0.83 at 14:07" can be joined to a specific model version in one query. It is also the cheapest smoke test for the runtime: an alerting rule that hits /model every minute and pages if the version changes unexpectedly catches most bad promotions.

Bake vs. pull, revisited​

Course 20 module 6 debated baking the model into the image against pulling it from a registry at startup. The two options translate directly into the MODEL_URI value:

  • Baked: MODEL_URI=file:/opt/model — startup reads a file, is offline-safe, and takes tens of milliseconds. Rollback is a redeploy of the previous image.
  • Pulled: MODEL_URI=models:/churn-classifier@production — startup calls MLflow, is coupled to registry availability, and takes seconds. Rollback moves the alias and restarts pods.

Whichever you pick, the lifespan shape above is identical. The choice is a deployment decision, not an application one, which is why keeping the URI in an environment variable — never in the source — is the pattern to follow.

Loading in a global module-level statement is not "loading at startup"

Writing model = mlflow.pyfunc.load_model(...) at the top of main.py also loads the model once. It looks equivalent, and it works in production, but it makes the tests painful: importing the module in a unit test now requires MLflow to be reachable. lifespan runs only when the ASGI server starts, so tests that import your handlers do not pay the load cost.

In summary​

  • Load the model once, in a lifespan context manager; never inside a request handler.
  • Read the MODEL_URI from an environment variable, resolve any alias to a concrete version at startup, and store both.
  • Let load failures abort the process so the orchestrator marks the pod NotReady — fail fast, not on the first user request.
  • Expose a small /model endpoint that returns URI, version and service version; join it with per-request logs (module 8).

Next module: single and batch prediction — one request for one subscriber, and one request for a vectorized batch, sharing the same preprocessing and the same model reference from this module.