Building Production AI Pipelines in Python, Step by Step

A production AI pipeline in Python is a chain of small, separately runnable steps (ingest, validate, build features, train, evaluate, register, serve, monitor) with an enforced contract at every boundary between them. The model is the smallest part. What makes the pipeline production-grade is that bad data fails loudly at ingestion, each feature is computed by one function in both training and serving, every model traces back to the exact data and code that produced it, and nothing reaches users without passing an automated evaluation gate.

This guide builds that pipeline end to end for one concrete case, a transaction fraud model, with the code for each stage and the failure each stage prevents.

Key takeaways

• Validate at the boundary. A strict schema on every table that enters the pipeline turns silent upstream changes into loud, early failures.
• Define features once. One function, imported by training and serving, with the decision time passed in. Two implementations of one feature, slowly drifting apart, cause some of the hardest ML bugs to find.
• Split by time, not at random. Random splits on event data leak the future into training and inflate offline metrics.
• Gate every promotion. A candidate must beat the current model on the same holdout and on the slices you care about, or it does not ship.
• Pin model versions in serving. A moving registry alias lets a registry click change production behavior without a deploy.
• Monitor inputs first. Labels arrive late; input and score distributions shift today.

The shape of a production pipeline

The stages are familiar. The contracts between them are what most notebook-to-production efforts skip.

Project layout

Keep the pipeline in one installable package. The training job, the serving image and the tests all import the same code, which is the cheapest way to make sure there is only one implementation of anything.

fraud-pipeline/
  pyproject.toml
  src/fraud/
    schemas.py      # data contracts for every table that crosses a boundary
    features.py     # the only place features are defined
    train.py
    evaluate.py
    serve.py
    flows.py        # orchestration
  tests/
    test_schemas.py
    test_features.py
    test_train_smoke.py

Pin dependencies with a lock file and build the serving image from the same lock. A model pickled under one scikit-learn version and loaded under another is a classic way to lose a Friday.

Step 1: ingest and validate at the boundary

Validate before anything is transformed, and fail rather than impute something plausible. pandera expresses the contract as a class, which makes it reviewable in a pull request like any other code.

"""src/fraud/schemas.py"""
import pandas as pd
import pandera.pandas as pa
from pandera.typing import Series
 
 
class Transactions(pa.DataFrameModel):
    transaction_id: Series[str] = pa.Field(unique=True)
    account_id: Series[str] = pa.Field(nullable=False)
    amount: Series[float] = pa.Field(ge=0, le=1_000_000)
    currency: Series[str] = pa.Field(isin=["USD", "EUR", "GBP", "INR"])
    created_at: Series[pd.Timestamp] = pa.Field(nullable=False)
 
    class Config:
        strict = True   # an unexpected column is an error, not a surprise
        coerce = True
 
 
def load_transactions(path: str) -> pd.DataFrame:
    df = pd.read_parquet(path)
    try:
        return Transactions.validate(df, lazy=True)
    except pa.errors.SchemaErrors as err:
        # lazy=True collects every failing check, so one run shows the whole problem
        err.failure_cases.to_csv("validation_failures.csv", index=False)
        raise

strict = True is the setting that earns its keep. Without it, an upstream team adding or renaming a column is invisible to you. With it, the pipeline fails on the first run after the change, while the cause is still obvious.

Schema checks catch shape problems. They do not catch a feed that delivered yesterday's file twice, or half a day of data. Add batch-level checks for volume and freshness:

def check_batch(df: pd.DataFrame, previous_rows: int, now: pd.Timestamp) -> None:
    if previous_rows and not (0.5 * previous_rows <= len(df) <= 2 * previous_rows):
        raise ValueError(f"row count {len(df)} vs {previous_rows} on the last run")
    lag = now - df["created_at"].max()
    if lag > pd.Timedelta(hours=6):
        raise ValueError(f"newest row is {lag} old; the feed may be stuck")

The bounds are examples; set yours from your own feed's history.

Two storage rules round this out. Keep raw data immutable, partitioned by ingest date, and write every transform to a new location rather than updating in place. Then record a data version (a Delta or Iceberg snapshot ID, a DVC hash, or simply an immutable path) for every dataset a training run reads.

Step 2: define features once, point-in-time correct

One of the most expensive bugs in production ML is training-serving skew: training computes a feature one way, serving computes it slightly differently, and the model quietly degrades while every dashboard stays green. Its sibling is leakage, where a training feature includes information that was not available at decision time.

Both have the same fix. Write each feature once, pass the decision time in, and prove that the batch and online executions agree.

"""src/fraud/features.py"""
from datetime import timedelta
 
import pandas as pd
 
WINDOW = timedelta(hours=24)
 
 
def velocity_24h(history: pd.DataFrame, as_of: pd.Timestamp) -> int:
    """Transactions for one account in the 24 hours strictly before `as_of`.
 
    `as_of` is passed in, never read from the clock, so a backfill for last
    March and a live request today run exactly the same logic.
    """
    in_window = (history["created_at"] >= as_of - WINDOW) & (history["created_at"] < as_of)
    return int(in_window.sum())
 
 
def velocity_24h_batch(txns: pd.DataFrame) -> pd.Series:
    """Vectorized version for training: one value per transaction row."""
    txns = txns.sort_values(["account_id", "created_at"])
    counts = (
        txns.set_index("created_at")
        .groupby("account_id")["amount"]
        .rolling(WINDOW, closed="left")   # excludes the current transaction
        .count()
        .fillna(0)
        .astype(int)
    )
    return pd.Series(counts.to_numpy(), index=txns.index, name="velocity_24h")

A function that calls datetime.now() internally cannot be backfilled correctly, and the skew it causes is very hard to spot after the fact. The test that keeps the two versions honest is short:

"""tests/test_features.py"""
from fraud.features import velocity_24h, velocity_24h_batch
 
 
def test_batch_matches_online(sample_txns):
    batch = velocity_24h_batch(sample_txns)
    for idx, row in sample_txns.sample(200, random_state=0).iterrows():
        history = sample_txns[sample_txns["account_id"] == row["account_id"]]
        assert batch[idx] == velocity_24h(history, row["created_at"])

When several models share features, or a feature has to be fresh within seconds, this pattern grows into a feature store: the same idea (one definition, point-in-time joins for training, low-latency reads for serving) with a registry and storage around it. Start with the shared package. Adopt the store when the package stops being enough.

Step 3: train reproducibly

Two decisions here matter more than the algorithm.

Split by time. Fraud, churn, demand and almost every business dataset are time-ordered. A random split puts next month's patterns in the training set and produces offline metrics that production will never match. Train on the past and validate on a later window. Drop rows too recent for their labels to have settled: a chargeback can arrive weeks after the transaction, so the newest rows are labeled "not fraud" when some of them are. A short gap between the training and validation windows also stops one fraud ring or campaign from straddling the split.

Tie the run to its inputs. Log the data version and code commit with the hyperparameters. A run you cannot connect to its exact training data is not reproducible, however carefully you logged the learning rate.

"""src/fraud/train.py"""
import mlflow
import mlflow.sklearn
import pandas as pd
from sklearn.ensemble import HistGradientBoostingClassifier
from sklearn.metrics import average_precision_score
 
FEATURES = ["amount", "velocity_24h", "account_age_days"]
LABEL_MATURITY = pd.Timedelta(days=30)   # chargebacks can take weeks to arrive
GAP = pd.Timedelta(days=7)
 
 
def train_candidate(df: pd.DataFrame, as_of: pd.Timestamp, cutoff: pd.Timestamp,
                    data_version: str, code_sha: str) -> str:
    mature = df[df["created_at"] < as_of - LABEL_MATURITY]    # labels old enough to trust
    train_df = mature[mature["created_at"] < cutoff - GAP]
    valid_df = mature[mature["created_at"] >= cutoff]
 
    params = {"max_iter": 300, "learning_rate": 0.05, "max_leaf_nodes": 31}
    mlflow.set_experiment("txn-fraud")
    with mlflow.start_run():
        mlflow.set_tags({"data_version": data_version, "code_sha": code_sha,
                         "as_of": as_of.isoformat(), "cutoff": cutoff.isoformat()})
        mlflow.log_params(params)
 
        model = HistGradientBoostingClassifier(**params, random_state=42)
        model.fit(train_df[FEATURES], train_df["is_fraud"])
 
        scores = model.predict_proba(valid_df[FEATURES])[:, 1]
        # PR AUC, not accuracy: with rare positives, accuracy rewards predicting "never fraud"
        mlflow.log_metric("valid_pr_auc", average_precision_score(valid_df["is_fraud"], scores))
 
        info = mlflow.sklearn.log_model(model, name="model",
                                        input_example=valid_df[FEATURES].head(5))
        return info.model_uri

as_of and cutoff are passed in for the same reason as in the feature code: a rerun of last month's training produces last month's model. Both go into the run tags, where you will want them when someone asks why last quarter's model behaved differently.

Step 4: put an evaluation gate in front of the registry

"The training job succeeded" means the code ran. It says nothing about whether the new model is better. The gate compares the candidate with the current champion on the same holdout, overall and on the slices where a regression would hurt, and fails closed.

"""src/fraud/evaluate.py"""
from sklearn.metrics import average_precision_score
 
from fraud.train import FEATURES
 
 
def pr_auc(model, frame):
    return average_precision_score(frame["is_fraud"], model.predict_proba(frame[FEATURES])[:, 1])
 
 
def gate(candidate, champion, holdout, slices, tolerance=0.005) -> list[str]:
    failures = []
    cand, champ = pr_auc(candidate, holdout), pr_auc(champion, holdout)
    if cand < champ - tolerance:
        failures.append(f"overall PR AUC {cand:.4f} vs champion {champ:.4f}")
 
    for name, mask in slices.items():          # e.g. "new_accounts", "currency_INR"
        part = holdout[mask]
        if part["is_fraud"].nunique() < 2:
            failures.append(f"slice {name} has no positives; cannot evaluate")
            continue
        c, p = pr_auc(candidate, part), pr_auc(champion, part)
        if c < p - tolerance:
            failures.append(f"slice {name}: {c:.4f} vs champion {p:.4f}")
    return failures

Add the checks a metric cannot express: the serialized model loads in the serving image, p99 inference latency fits the budget, and the score distribution on the holdout is not degenerate (everything near zero is a common symptom of a broken feature). Any failure stops the flow and pages a person. An empty slice is treated as a failure on purpose; a gate that silently skips what it cannot measure is not a gate.

Step 5: register and promote deliberately

MLflow deprecated registry stages ("Staging", "Production") in favor of aliases and tags. Use an alias to mark the candidate, but have serving load an exact version, so changing what production runs is a deploy with a diff and a rollback, not a click.

import mlflow
from mlflow import MlflowClient
 
client = MlflowClient()
mv = mlflow.register_model(model_uri, "txn-fraud")
client.set_registered_model_alias("txn-fraud", "challenger", mv.version)

Promotion is then a config change reviewed like code: the serving deployment's MODEL_URI moves to models:/txn-fraud/<version>, and rolling back means reverting that line.

Step 6: serve with the same contracts

Pick the serving mode from the decision, not from fashion.

For the synchronous case, the service validates input as strictly as the pipeline did, builds features with the shared module, loads a pinned version once at startup, and logs what it used.

"""src/fraud/serve.py"""
import logging
import os
from contextlib import asynccontextmanager
 
import mlflow.sklearn
from fastapi import FastAPI
from pydantic import BaseModel, Field
 
from fraud.features import build_online_features   # same package as training
from fraud.train import FEATURES
 
MODEL_URI = os.environ["MODEL_URI"]                # an exact version, never a moving alias
log = logging.getLogger("fraud.serve")
state = {}
 
 
@asynccontextmanager
async def lifespan(app: FastAPI):
    state["model"] = mlflow.sklearn.load_model(MODEL_URI)
    yield
 
 
app = FastAPI(lifespan=lifespan)
 
 
class ScoreRequest(BaseModel):
    transaction_id: str
    account_id: str
    amount: float = Field(ge=0, le=1_000_000)
    currency: str = Field(pattern="^(USD|EUR|GBP|INR)$")
 
 
@app.post("/score")
def score(req: ScoreRequest) -> dict:
    row = build_online_features(req)                # one-row DataFrame
    proba = float(state["model"].predict_proba(row[FEATURES])[0, 1])
    log.info("scored", extra={"transaction_id": req.transaction_id, "model_uri": MODEL_URI,
                              "features": row[FEATURES].iloc[0].to_dict(), "score": proba})
    return {"fraud_probability": proba, "model": MODEL_URI}

The prediction log is not optional. It is the input to monitoring, the evidence for skew investigations, and the training data for the next model once labels arrive. Define it as a schema like any other table.

Step 7: monitor inputs, scores and delayed labels

Model quality needs ground truth, and ground truth often arrives days or weeks late. The first signals you can act on are shifts in the inputs and in the score distribution. The population stability index (PSI) is a simple, explainable drift measure you can compute without a framework:

import numpy as np
 
 
def psi(reference: np.ndarray, current: np.ndarray, bins: int = 10) -> float:
    """Population stability index for one numeric feature or for model scores."""
    edges = np.unique(np.quantile(reference, np.linspace(0, 1, bins + 1)))
    edges[0], edges[-1] = -np.inf, np.inf
    ref = np.histogram(reference, edges)[0] / len(reference)
    cur = np.histogram(current, edges)[0] / len(current)
    ref, cur = np.clip(ref, 1e-6, None), np.clip(cur, 1e-6, None)
    return float(np.sum((cur - ref) * np.log(cur / ref)))

A common credit-risk rule of thumb reads PSI below 0.1 as stable and above 0.25 as a significant shift. Treat those as starting points and calibrate against your own history, because a feature with a weekly cycle will look like drift every Monday. Evidently and NannyML package this and more (NannyML can estimate performance before labels arrive) if you would rather not maintain it.

Monitor in layers, and alert only on what someone would act on:

The last row catches the quietest failure of all: the retraining job stopped running a month ago and nothing errored, because nothing ran. Alert on absence. For the plumbing behind these signals (structured logs, metrics, traces), see our guide to observability in production.

Orchestrate the whole thing

The orchestrator runs the steps in order, retries what is safe to retry, and records what happened. In Prefect, the flow is ordinary Python:

"""src/fraud/flows.py"""
from prefect import flow, task
 
 
@task(retries=2, retry_delay_seconds=300)
def ingest(day: str) -> str:
    """Copy the day's raw files to an immutable path; return that path."""
    ...
 
 
@task
def build_training_set(raw_path: str) -> str:
    """Validate, build features, write a versioned dataset; return its version."""
    ...
 
 
@task
def train(dataset_version: str) -> str:
    ...
 
 
@task
def evaluate_and_register(model_uri: str) -> None:
    """Run the gate; raise if it fails so the flow run is marked failed."""
    ...
 
 
@flow(name="fraud-weekly-retrain")
def retrain(day: str) -> None:
    raw = ingest(day)
    dataset = build_training_set(raw)
    model_uri = train(dataset)
    evaluate_and_register(model_uri)

Retries belong only on idempotent tasks. Every task writes to a new, run-keyed location, so a retry or a rerun of last Tuesday produces the same result instead of a duplicate. The same structure maps directly onto Airflow DAGs or Dagster assets. For how this layer sits inside a wider data platform, see the modern data engineering stack.

What goes wrong in production

These are the failures that wake people up. Almost none of them are about the model.

• An upstream schema changed and nothing validated it. A renamed column arrived as nulls, a default filled them with zeros, and the model scored every request as low risk. Strict schemas exist for this.
• Training and serving computed a feature differently. Offline metrics never matched online behavior, and it took weeks to find because each implementation looked correct on its own.
• The retrain succeeded on corrupted data and promoted itself. Automatic promotion without a gate turns one bad batch into a production incident.
• The split leaked the future. A random split, or a feature computed with data from after the decision time, produced a model that was excellent offline and mediocre live.
• A registry alias moved and production changed with no deploy. Nobody could say when behavior changed or why. Pin versions.
• The scheduled job stopped running. The model aged quietly while every dashboard stayed green. Alert on model age and missing runs.
• A library upgrade broke model loading. The serving image resolved a newer dependency than the training environment. Build both from one lock file.
• Nobody could reproduce the model in production. No data version, no code commit, no way to answer an auditor or debug a regression.

Frequently Asked Questions

What is an AI pipeline in Python?

An AI pipeline is the chain of steps that turns raw data into predictions users rely on: ingestion, validation, feature computation, training, evaluation, registration, serving and monitoring. In Python each step is usually a function or task in a shared package, run by an orchestrator such as Prefect, Airflow or Dagster. What makes it a production pipeline rather than a notebook is that every step can be rerun on its own and every boundary between steps is checked in code.

Which Python tools should I use for a machine learning pipeline?

A dependable default is pandera or Great Expectations for data validation, pandas or Polars for transforms, scikit-learn, LightGBM or PyTorch for models, MLflow for experiment tracking and the model registry, FastAPI for online serving, and Prefect, Airflow or Dagster for orchestration. The specific tools matter less than the contracts between them. Pick boring, well-documented tools your team already knows and spend the saved time on validation and evaluation.

How do you prevent training-serving skew?

Define each feature exactly once, in a package that both the training job and the serving code import, and pass the decision time in explicitly instead of reading the clock inside the function. Add a test that computes features through the batch path and the online path for the same rows and asserts they match. Log the features the service actually used, so you can compare them with offline recomputation after the fact.

Should I use Airflow, Prefect or Dagster for ML pipelines?

All three can run a production ML pipeline. Airflow is the most established and fits teams that already operate it for data engineering, Prefect has the lightest path from plain Python functions to scheduled flows, and Dagster is built around data assets and lineage, which suits feature and dataset heavy work. Choose the one your team will actually operate well; the pipeline design matters far more than the scheduler.

How often should a production model be retrained?

As often as the data changes in ways that hurt the model, which you learn by monitoring rather than by guessing. Many teams start with a fixed schedule, such as weekly, and add drift-triggered runs once they trust their monitoring. Whatever the trigger, a retrained model should pass the same evaluation gate as any other candidate before it serves traffic.

Do the same practices apply to LLM pipelines?

Mostly, yes. You still validate inputs at the boundary, version prompts and retrieval code like any other code, gate releases on an evaluation set, and log inputs, outputs and the model version for every call. The differences are that evaluation sets are built from task examples instead of labeled tables, and cost and latency per call need their own monitoring.

Need help getting a model into production?

Most of our AI work starts where a notebook or a demo stops: the model works, and now it has to survive real data, real traffic and a retraining schedule. Senior engineers build the pipeline, the gates and the monitoring, and if we built it, we are on the pager for it.

• AI & Machine Learning practice
• POC-to-Production Engineering
• Talk to an engineer