Amit chakraborty

Posted on Oct 15 Originally published at amitchakraborty.dev

MLOps: models that survive production · 3. Features: stores, and the simpler thing that usually works

#tutorial #aiengineering #mlops

MLOps: models that survive production · Chapter 3 of 12 · AI Engineering · new chapter every Thursday night

By the end of this chapter: Avoid training and serving skew without adopting a platform you do not need.

The problem

Say you are Maya, a data‑science‑engineer on a two‑person ML team at a fintech startup. You have just finished a churn‑prediction model in a Jupyter notebook using a CSV export from the data‑warehouse. The notebook runs end‑to‑end: you read the CSV, apply a handful of pandas transformations –‑ one‑hot‑encode a “plan_type” column, fill missing “balance” values with the median, and create a “days_since_last_login” feature from a timestamp. After a quick train‑test split, a LightGBM model hits 0.92 AUC. You push the model to the serving repo, wrap the same preprocessing code in a Flask endpoint, and ship it behind the API gateway.

Two days later, the on‑call engineer receives an alert: the new endpoint is returning a 500 error for 12 % of requests. The logs show a KeyError: 'plan_type' inside the preprocessing step. A quick look at the data warehouse reveals that the “plan_type” column has been renamed to “subscription_type” overnight as part of a schema migration. The model now receives a missing feature, the preprocessing crashes, and the service is down for a chunk of users. The outage costs the team an estimated £8 k in lost transactions and a week of firefighting.

Maya’s notebook‑centric workflow gave her a perfect model on static data, but the production pipeline suffered from training‑serving feature skew – the features used at training time differed from those at serving time. The root cause was the lack of a systematic way to store, version, and reuse features across environments, not a problem with the model itself.

The goal of this chapter is to give you a lightweight, production‑ready pattern that eliminates this skew without pulling in a heavyweight feature‑store platform. By the end you will be able to:

Define a deterministic feature‑generation library that can be imported both in notebooks and in the serving code.
Store feature artefacts (e.g., Parquet files) with explicit version identifiers.
Load the exact same feature version at training and at inference time, guaranteeing feature parity.

Before you start

Prerequisite	Version / command
Python	3.10
pandas	2.2.0 (`pip install pandas==2.2.0`)
pyarrow (for Parquet)	15.0.0 (`pip install pyarrow==15.0.0`)
scikit‑learn	1.5.0 (`pip install scikit-learn==1.5.0`)
LightGBM	4.5.0 (`pip install lightgbm==4.5.0`)
Git	any recent version (used for version control of feature code)

Verify the environment – run the following one‑liner; you should see the versions printed without error:

python - <<'PY'
import pandas, pyarrow, sklearn, lightgbm, sys
print(f"pandas {pandas.__version__}")
print(f"pyarrow {pyarrow.__version__}")
print(f"sklearn {sklearn.__version__}")
print(f"lightgbm {lightgbm.__version__}")
PY

Expected output (exact numbers may differ if you installed a newer patch, but they must be ≥ the versions above):

pandas 2.2.0
pyarrow 15.0.0
sklearn 1.5.0
lightgbm 4.5.0

If any import fails, reinstall the offending package before proceeding.

How to keep feature generation deterministic?

1. Why a shared library beats copy‑paste

In many organisations the first instinct is to copy the preprocessing cells from the notebook into the serving script. This creates two independent code bases that diverge as soon as a developer tweaks a column name or a default value. The single source of truth principle tells us to extract the logic into a reusable Python module that both the notebook and the service import.

What to do: Create a package features/ with an __init__.py exposing a build_features(df: pd.DataFrame) -> pd.DataFrame function. Keep the function pure – no external state, no random seeds (unless you explicitly set them). This guarantees that given the same input DataFrame, the output is identical regardless of where the function runs.

How to confirm: Run the function on a small CSV twice – once from the notebook, once from a plain Python script – and compare the hashes of the resulting DataFrames.

import hashlib, pandas as pd, features

df = pd.read_csv('sample.csv')
feat = features.build_features(df)

hash1 = hashlib.sha256(feat.to_json().encode()).hexdigest()
print(hash1)

Run the same snippet from the notebook; the printed hash must match exactly.

2. How to version the feature artefacts

Even with a shared library, you still need to guarantee that the data used to train the model matches the data the model will see at inference time. The simplest approach is to materialise the feature set as a Parquet file and tag it with a version identifier derived from the Git commit SHA of the feature library.

What to do: After you run build_features on the full training dataset, write the result to s3://my‑ml‑bucket/features/v{git_sha}.parquet. Store the same git_sha in the model artefact metadata (e.g., a JSON file alongside the model pickle). At serving time, the inference code reads the version from the model metadata and loads the corresponding Parquet file for any offline feature look‑ups (e.g., for batch scoring) or validates that the online feature extraction matches the schema.

How to confirm: After training, inspect the model metadata:

{
  "model_path": "models/churn_v1.pkl",
  "feature_version": "a1b2c3d4"
}

Then, in the serving container, run:

aws s3 ls s3://my-ml-bucket/features/ | grep a1b2c3d4

You should see exactly one file: features/a1b2c3d4.parquet.

3. How to guard against schema drift

A schema change (e.g., column rename) will raise a KeyError at runtime if the feature library expects a missing column. To catch this early, add a validation step that asserts the input DataFrame contains all required columns before any transformation.

What to do: In features.build_features, compute the set difference between required_columns and df.columns. If any are missing, raise a custom FeatureSchemaError with a clear message.

How to confirm: Run the function on a deliberately broken DataFrame:

import pandas as pd, features
df = pd.DataFrame({'balance': [100, 200]})  # missing 'plan_type'
try:
    features.build_features(df)
except features.FeatureSchemaError as e:
    print(e)

Expected output:

FeatureSchemaError: Missing required columns: {'plan_type', 'last_login_ts'}

Deploy this validation into both training and serving pipelines; the error will surface during CI tests before code reaches production.

Build it

Below is the end‑to‑end recipe that turns the abstract ideas above into a concrete, runnable artefact. Follow each step, verify the indicated output, and you will have a reproducible feature pipeline that eliminates training‑serving skew.

Step 1 – Initialise a feature package

mkdir -p ml_project/features
touch ml_project/features/__init__.py

Edit ml_project/features/__init__.py:

import pandas as pd
from typing import List

class FeatureSchemaError(RuntimeError):
    """Raised when input data does not contain the expected columns."""

REQUIRED_COLUMNS = [
    "customer_id",
    "plan_type",
    "balance",
    "last_login_ts",
    "signup_ts"
]

def _validate(df: pd.DataFrame) -> None:
    missing = set(REQUIRED_COLUMNS) - set(df.columns)
    if missing:
        raise FeatureSchemaError(f"Missing required columns: {missing}")

def _one_hot(df: pd.DataFrame, column: str) -> pd.DataFrame:
    dummies = pd.get_dummies(df[column], prefix=column)
    return pd.concat([df.drop(columns=[column]), dummies], axis=1)

def _days_since(series: pd.Series, reference: pd.Timestamp) -> pd.Series:
    return (reference - pd.to_datetime(series)).dt.days

def build_features(df: pd.DataFrame) -> pd.DataFrame:
    """
    Deterministic feature generation.
    Returns a new DataFrame with engineered columns.
    """
    _validate(df)

    # 1. One‑hot encode plan_type
    df = _one_hot(df, "plan_type")

    # 2. Median imputation for balance
    median_balance = df["balance"].median()
    df["balance"] = df["balance"].fillna(median_balance)

    # 3. Days since last login (relative to a fixed reference date)
    reference_date = pd.Timestamp("2024-01-01")
    df["days_since_last_login"] = _days_since(df["last_login_ts"], reference_date)

    # 4. Drop raw timestamps – they are no longer needed
    df = df.drop(columns=["last_login_ts", "signup_ts"])

    return df

Why: The module contains a single public entry point, deterministic transformations, and a strict schema guard.

Confirm: Run a quick sanity check.

python - <<'PY'
import pandas as pd, ml_project.features as ft
df = pd.DataFrame({
    "customer_id": [1,2],
    "plan_type": ["basic","premium"],
    "balance": [100.0, None],
    "last_login_ts": ["2023-12-20","2023-12-25"],
    "signup_ts": ["2022-01-15","2022-03-10"]
})
out = ft.build_features(df)
print(out.head())
PY

Expected output (order of columns may vary but all listed columns must appear, and no KeyError should be raised):

   customer_id  balance  plan_type_basic  plan_type_premium  days_since_last_login
0            1    100.0                1                  0                     12
1            2    100.0                0                  1                      7

Step 2 – Version the feature artefact

First, obtain the current Git SHA of the repository (assume you have a Git repo at the project root).

git init  # if not already a repo
git add .
git commit -m "Initial commit with feature library"
FEATURE_SHA=$(git rev-parse --short HEAD)
echo $FEATURE_SHA   # e.g. a1b2c3d4

Now, generate the full training feature set and write it to S3 (or local storage for the demo).

python - <<'PY'
import pandas as pd, ml_project.features as ft, os, json, subprocess, sys
# Load raw training data – replace with your actual source
raw = pd.read_csv("data/training_raw.csv")
features = ft.build_features(raw)

# Write Parquet locally (replace with s3 path in production)
version = subprocess.check_output(["git", "rev-parse", "--short", "HEAD"]).decode().strip()
out_path = f"features/v{version}.parquet"
features.to_parquet(out_path, engine="pyarrow")
print(f"Written {out_path}")

# Persist model metadata
metadata = {
    "model_path": "models/churn_v1.pkl",
    "feature_version": version
}
with open("models/churn_v1_metadata.json", "w") as f:
    json.dump(metadata, f, indent=2)
print("Metadata written")
PY

Confirm:

ls features/

You should see a file like v a1b2c3d4.parquet. Also inspect the metadata:

cat models/churn_v1_metadata.json

Expected JSON (SHA will match your repo):

{
  "model_path": "models/churn_v1.pkl",
  "feature_version": "a1b2c3d4"
}

Step 3 – Load the same feature version at inference time

Create a tiny Flask inference service that reads the model metadata, loads the matching feature version, and applies the same build_features function to incoming JSON payloads.

mkdir -p ml_project/serve
touch ml_project/serve/app.py

Edit ml_project/serve/app.py:

import json
import pandas as pd
import joblib
from flask import Flask, request, jsonify
import ml_project.features as ft
import subprocess

app = Flask(__name__)

# Load model and metadata once at startup
model = joblib.load("models/churn_v1.pkl")
with open("models/churn_v1_metadata.json") as f:
    meta = json.load(f)
feature_version = meta["feature_version"]

# Pre‑load the feature artefact (useful for batch scoring; for online we just validate schema)
feature_path = f"features/v{feature_version}.parquet"
stored_features = pd.read_parquet(feature_path)

@app.route("/predict", methods=["POST"])
def predict():
    # Expect a JSON list of records
    records = request.get_json()
    df_raw = pd.DataFrame.from_records(records)

    # Validate schema early
    try:
        ft._validate(df_raw)
    except ft.FeatureSchemaError as e:
        return jsonify({"error": str(e)}), 400

    # Build features deterministically
    df_feat = ft.build_features(df_raw)

    # Align columns with stored feature set (ensures ordering & missing columns are caught)
    missing = set(stored_features.columns) - set(df_feat.columns)
    if missing:
        return jsonify({"error": f"Missing engineered columns: {missing}"}), 400
    df_feat = df_feat[stored_features.columns]

    preds = model.predict_proba(df_feat)[:, 1].tolist()
    return jsonify({"predictions": preds})

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=8080)

Why: The service reads the same feature_version that the model was trained with, validates the raw input against the required schema, and forces column ordering to match the stored artefact. Any deviation (e.g., a renamed column) will be caught with a 400 response rather than a 500 crash.

Confirm: Start the server locally.

python ml_project/serve/app.py

In another terminal, send a valid request:

curl -X POST http://127.0.0.1:8080/predict \
  -H "Content-Type: application/json" \
  -d '[{"customer_id":1,"plan_type":"basic","balance":120.0,"last_login_ts":"2023-12-20","signup_ts":"2022-01-15"}]'

Expected JSON response (probability will vary):

{
  "predictions": [0.73]
}

Now send a request with a missing column:

curl -X POST http://127.0.0.1:8080/predict \
  -H "Content-Type: application/json" \
  -d '[{"customer_id":1,"balance":120.0,"last_login_ts":"2023-12-20","signup_ts":"2022-01-15"}]'

Expected output:

{
  "error": "FeatureSchemaError: Missing required columns: {'plan_type'}"
}

The service gracefully rejects the malformed payload instead of raising an internal server error.

When this breaks

Symptom (what you see)	Cause	Fix
`KeyError: 'plan_type'` inside `build_features`	Raw input missing a required column (schema drift)	Ensure the data‑warehouse export includes the column, or update `REQUIRED_COLUMNS` and re‑train. Add a migration test that checks the warehouse schema against `REQUIRED_COLUMNS`.
`ValueError: cannot convert string to float: 'N/A'` during median imputation	Unexpected non‑numeric sentinel values in a numeric column	Extend `build_features` to coerce non‑numeric entries to `np.nan` before imputation: `df['balance'] = pd.to_numeric(df['balance'], errors='coerce')`.
`FileNotFoundError: [Errno 2] No such file or directory: 'features/vdeadbeef.parquet'` at inference start	Model metadata points to a feature version that was never materialised (e.g., CI skipped the artefact step)	Add a CI gate that asserts the artefact exists: `aws s3 ls s3://my-ml-bucket/features/v${FEATURE_SHA}.parquet`.
`AssertionError: columns mismatch` when aligning `df_feat` with `stored_features`	New engineered column added after training but before deployment (e.g., you added a new one‑hot level)	Increment the feature version (commit the change, re‑run the training pipeline, produce a new artefact). Deploy the new model together with the new feature version atomically.
Latency spikes > 200 ms on `/predict`	Online feature extraction now performs a heavy join or external API call that was not present during training	Keep online feature extraction lightweight; offload heavy joins to an offline batch pipeline and store the result in the versioned Parquet artefact.

Each fix restores the invariant: the exact same deterministic transformation, applied to the same schema, produces identical columns for training and serving.

What it costs

Storage – Every feature version is a full Parquet copy of the training data (often tens of GB). Retaining many versions can quickly inflate S3 costs. Mitigation: prune versions older than the last 3 releases, or use delta‑lake techniques to store only diffs (beyond this chapter’s scope).
Operational overhead – You must enforce a CI step that checks the feature version exists and that the schema validation passes. This adds a few minutes to each pipeline run.
Latency – The online path still runs the same deterministic code as offline. If a transformation is computationally heavy, it will affect request latency. The trade‑off is to keep the online path minimal; heavy feature engineering belongs in the offline artefact.
Developer discipline – All team members must import ml_project.features everywhere. A stray copy‑paste will silently re‑introduce skew. Code‑review checklists should flag any direct pandas manipulation outside the library.

When the project grows beyond a handful of models, you may consider a dedicated feature‑store (e.g., Feast) to manage lineage, TTL, and online serving caches. Until then, the pattern above provides a low‑maintenance, auditable bridge between notebook experiments and production inference.

In the interview

Typical interview prompt

“Explain how you would prevent training‑serving feature skew in a production ML system without using a managed feature store.”

Weak answer – “I would just copy the preprocessing code from my notebook into the service and make sure the column names match.”

Why it’s weak: It shows no awareness of versioning, schema validation, or the need for a single source of truth. It also ignores the possibility of future schema changes.

Strong answer – “I would extract all preprocessing into a pure Python module that both notebooks and the service import. I would version the generated feature set by committing the feature‑library SHA and persisting the Parquet artefact under that version. The model metadata would store the same version identifier, and the service would load the exact artefact at start‑up, validating the raw input schema before transformation. This guarantees deterministic parity and catches schema drift early. The trade‑off is extra storage for each version and a CI gate that checks artefact existence, but it avoids costly 500 errors in production.”

What it names: single source of truth, deterministic function, version‑tagged artefacts, schema validation, storage/CI trade‑off, failure mode (schema drift).

Senior‑level extension – A senior candidate would also discuss lineage tracking (e.g., storing the Git commit hash, data‑snapshot hash), back‑fill strategies for historic data when a new feature is added, and automated rollback (loading the previous feature version if the new one fails health checks). They might also mention feature‑store evaluation criteria and why the lightweight approach is preferable for a small team.

Director/manager angle – The interviewer will probe the business impact: “What cost does this approach incur, and how would you justify it to leadership?” A good answer quantifies storage (e.g., £0.023 per GB‑month on S3), explains the risk reduction (e.g., avoiding £10 k outages), and outlines a governance process (retention policy, audit logs).

Follow‑up to separate theory from practice – “Can you walk me through a recent incident where a schema change broke your model, and how the versioned feature artefact would have prevented it?” A candidate who has lived the incident will describe the exact error, the manual fix, and the post‑mortem action of adding validation and versioning.

Your tasks

Re‑implement the feature library – Clone the features/__init__.py from this chapter, but add a new engineered column: log_balance = np.log1p(df["balance"]). Update the schema validation accordingly and re‑train the model.

Done when the model metadata reflects the new Git SHA and the inference endpoint returns predictions without error.
Simulate a schema drift – Rename the column plan_type to subscription_type in the raw CSV, commit the change, and run the inference service without updating the feature library. Verify that the service returns a 400 error with the correct FeatureSchemaError message.

Done when the error message appears exactly as shown in the “When this breaks” table.
Implement a CI check – Write a Bash script ci/check_features.sh that (a) extracts the SHA from models/churn_v1_metadata.json, (b) confirms the corresponding Parquet file exists in features/, and (c) exits with non‑zero status if not. Run the script locally to see both success and failure paths.

Done when ./ci/check_features.sh prints “Feature artefact v… exists” for the current version and fails with “Missing feature artefact v…”.
Add a retention policy – Write a Python script maintenance/prune_features.py that lists all feature versions in features/, keeps the latest three, and deletes older ones (use os.remove). Run it in a test directory with dummy files and verify only the three newest remain.

Done when after execution only three versioned Parquet files are left.
Break‑and‑diagnose – Introduce a subtle bug: in build_features, change the median imputation line to df["balance"] = df["balance"].fillna(0). Retrain the model, redeploy, and then feed a payload where balance is missing. Observe the prediction shift (e.g., all predictions drop dramatically). Use the stored feature artefact to compare the distribution of balance before and after the change, and write a short paragraph explaining why the bug caused a performance regression.

Done when you can point to the mismatched median value in the artefact and explain the impact on model output.

Your tasks this week

Do the exercises above before the next chapter. Reading a tutorial and doing
one are different activities and only one of them changes what you can build.

Stuck on any of them? Say so — describe what you tried and what happened:
tell me where you got stuck. I read every one, and the questions
that come back more than twice get answered in the next chapter.