Skip to main content

Module 4 — Single and batch prediction

The service now has a real model loaded once at startup. This module wires it to two shapes of traffic that every real ML API supports: one request for one subscriber, called by a dashboard or a live workflow, and one request for a batch of dozens to thousands of subscribers, called by a nightly job or a Streamlit CSV upload. Serving both cheaply, and with no divergence between them and the training pipeline, is what this module does.

Two routes, one model​

The single-row route is the direct evolution of module 3. The batch route accepts a list of subscribers and returns a list of predictions, aligned by position:

# app/main.py (continued)
from typing import List
from fastapi import FastAPI, HTTPException, status
from pydantic import BaseModel, Field, conlist

from app.schemas import Subscriber, Prediction
from app.state import MODEL_STATE

MAX_BATCH_SIZE = 500


class BatchIn(BaseModel):
    items: conlist(Subscriber, min_length=1, max_length=MAX_BATCH_SIZE)


class BatchOut(BaseModel):
    items: List[Prediction]
    model_version: str


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


@app.post("/predict", response_model=Prediction, tags=["prediction"])
def predict_one(subscriber: Subscriber) -> Prediction:
    return _score_one(subscriber)


@app.post("/predict/batch", response_model=BatchOut, tags=["prediction"])
def predict_batch(payload: BatchIn) -> BatchOut:
    # The single-row route reuses this same function, so vectorization
    # is the only difference between the two shapes.
    predictions = _score_many(payload.items)
    return BatchOut(items=predictions, model_version=MODEL_STATE["version"])

Two design notes deserve names. conlist(..., max_length=500) turns "the batch is too big" into a 422 at the boundary, before a single row is scored — no CPU is spent on a request the service will refuse. A single scoring function (_score_many, defined below) is the source of truth: the single-row route calls it with a list of one. That prevents the split-brain bug where the two routes silently disagree because one was updated and the other was not.

Vectorized scoring, not a Python loop​

The scoring function itself must be vectorized. Every serious ML library takes a matrix and returns a vector in roughly constant time overhead. A Python for loop that calls model.predict([row]) once per row is 10 to 100 times slower on a batch of 500 rows, for reasons that have nothing to do with the model:

# app/scoring.py
from typing import List
import pandas as pd

from app.schemas import Subscriber, Prediction
from app.state import MODEL_STATE
from app.preprocess import to_model_frame


def _score_many(items: List[Subscriber]) -> List[Prediction]:
    model = MODEL_STATE["model"]
    version = MODEL_STATE["version"]

    # 1. Build one DataFrame for the whole batch.
    frame: pd.DataFrame = to_model_frame([s.model_dump() for s in items])

    # 2. One model call.
    probas = model.predict(frame)          # shape (n,)
    probas = [float(p) for p in probas]    # JSON-serializable

    # 3. Assemble the aligned output.
    out: List[Prediction] = []
    for sub, p in zip(items, probas):
        bucket = "low" if p < 0.33 else "medium" if p < 0.66 else "high"
        out.append(Prediction(
            subscriber_id=sub.subscriber_id,
            churn_probability=round(p, 4),
            risk_bucket=bucket,
            model_version=version,
        ))
    return out


def _score_one(subscriber: Subscriber) -> Prediction:
    # The single-row route calls the batch function with a list of one.
    return _score_many([subscriber])[0]

The two callable-cost lines are to_model_frame(...) (preprocessing on a batch) and model.predict(frame) (one call). Everything else — building the response, resolving buckets — runs in native Python and is negligible for batches of a few hundred rows. Module 10 will measure exactly how this scales with batch size.

The single most dangerous bug: divergent preprocessing​

The model was trained on features that had a specific one-hot encoding, a specific set of columns in a specific order, a specific handling of missing values, and possibly a specific scaler. Any drift between how the training pipeline preprocessed a row and how the service preprocesses the same row produces a training/serving skew: the request is well-formed, the model returns a probability, and that probability is quietly wrong.

The fix is to reuse the same preprocessing code in both places. There are two viable patterns:

  • Preprocessing baked into the model artifact. A scikit-learn Pipeline that starts with a ColumnTransformer and ends with the classifier is saved as a single object. The service receives raw features and the pipeline handles the encoding. This is the pattern MLflow's pyfunc flavor already supports, and it is the default recommendation.
  • A shared package that both the training job and the serving image import. The training job calls to_model_frame before fit; the service calls to_model_frame before predict. The package is versioned like any other dependency and travels with the model in the lock file.

An anti-pattern to avoid at all costs is rewriting preprocessing in the service because "we do not want the whole scikit-learn dependency". The saving is small (tens of megabytes) and the cost of a drift bug — a silently miscalibrated model that ships for a week — is enormous.

Here is what to_model_frame looks like for the churn model:

# app/preprocess.py
import pandas as pd

TRAIN_COLUMNS = [
    "tenure_months", "monthly_charges", "total_charges", "is_senior",
    "contract_type_month_to_month", "contract_type_one_year", "contract_type_two_years",
    "payment_method_bank_transfer", "payment_method_credit_card", "payment_method_mailed_check",
]


def to_model_frame(rows: list[dict]) -> pd.DataFrame:
    df = pd.DataFrame(rows)
    df = pd.get_dummies(df, columns=["contract_type", "payment_method"], dtype=int)
    # Force every training-time column to exist, with 0 for missing categorical levels.
    for col in TRAIN_COLUMNS:
        if col not in df.columns:
            df[col] = 0
    return df[TRAIN_COLUMNS]  # exact column order matters for tree models

The reindex on the last line is the reason this code exists at all. pd.get_dummies produces different columns depending on which categorical levels appear in the input batch. Without the reindex, a batch of subscribers who all have contract_type="one_year" is missing the other two one-hot columns, and the model receives a matrix with the wrong shape or wrong order.

Capping batch size, and why​

MAX_BATCH_SIZE = 500 is not a decorative constant. Above a few hundred rows, three costs start to matter:

  • JSON serialization. Parsing a 5 MB body and serializing the response takes tens of milliseconds on its own. FastAPI uses orjson under the hood if available, which helps, but there is no free lunch.
  • Memory. A batch of 5 000 rows with 20 features and one-hot expansion is a few hundred kilobytes; a batch of 500 000 is a few hundred megabytes and can OOM a small container.
  • p99 latency. A batch route with no cap is a denial-of-service vector: any caller can hold a worker for minutes with a single call. Capping at 500 keeps a bad caller from monopolizing a worker.

For traffic that legitimately needs more than 500 rows, the pattern is a background job, covered in module 6, backed by a file upload rather than a synchronous request.

Testing the two shapes together​

The two-route design keeps tests short and precise. Below, both routes must agree on the same subscriber to the last digit:

# tests/test_predict.py
from fastapi.testclient import TestClient

from app.main import app

client = TestClient(app)


def test_single_and_batch_agree_on_one_subscriber() -> None:
    payload = {
        "subscriber_id": "sub-000123", "tenure_months": 12,
        "monthly_charges": 75.5, "total_charges": 900.0,
        "contract_type": "month_to_month", "payment_method": "credit_card",
        "is_senior": False,
    }

    r1 = client.post("/predict", json=payload)
    r2 = client.post("/predict/batch", json={"items": [payload]})

    assert r1.status_code == 200 and r2.status_code == 200
    assert r1.json()["churn_probability"] == r2.json()["items"][0]["churn_probability"]


def test_batch_over_limit_is_rejected() -> None:
    payload = {"items": [_dummy() for _ in range(501)]}
    r = client.post("/predict/batch", json=payload)
    assert r.status_code == 422

The first test is the anti-drift guard: if any refactor makes _score_one diverge from _score_many, CI turns red. The second test locks in the cap so a well-meaning teammate does not silently double it.

Return a stable ID with every prediction

Even when the caller provides subscriber_id, echo it in the response. It costs nothing, and it turns "your API returned bad numbers for subscriber sub-000123" into a debuggable ticket instead of guesswork about which row of a 300-item batch went wrong.

In summary​

  • Expose a single-row route and a batch route that share a single vectorized scoring function — the split-brain bug does not exist because there is nothing to diverge.
  • Vectorize: one model call per batch, not one per row. The overhead of a Python loop dwarfs the model's own cost for anything above a handful of rows.
  • Guarantee preprocessing consistency with training by baking the pipeline into the artifact or by sharing a versioned preprocessing package. Never rewrite preprocessing in the service.
  • Cap batches with conlist(max_length=...); route legitimately huge requests to a background job (module 6) instead of stretching the synchronous path.

Next module: error handling and status codes — turning the raw exceptions of this module into structured, actionable HTTP responses that never leak a Python traceback.