pytest (pip install torch pytest covers the essentials);Learn AI Series):In traditional software, CI/CD is a solved problem, more or less. You push code, a runner spins up, it runs your tests, it builds an artifact, and if everything is green it ships that artifact to production. The pieces are well understood and the tooling is mature. You can go read a hundred good tutorials on it. Machine learning takes that comfortable picture and quietly sets it on fire.
Why? Because the thing you ship in ML is not compiled code. It's a trained model -- a big bag of numbers that fell out of a training run which depended on data, hyperparameters, a random seed, a particular library version, and quite possibly the phase of the moon. You cannot rebuild it deterministically the way you rebuild a binary from source. And your "tests" are not just asserting that add(2, 2) == 4. They have to assert that the data looks sane, that the model performs above some bar, that it isn't worse than the one already serving users, and that it hasn't silently learned to be unfair. That's a lot more surface area than a normal test suite.
Back in episode #34 we made the jump from notebook to production, and in #119 we learned to track experiments so a training run is reproducible. Today we close the loop: we take all that discipline and we automate it, so that shipping a model is as boring and reliable as shipping a web app should be. Because here's the uncomfortable truth -- if your ML pipeline still depends on someone opening a notebook, running the cells top to bottom, eyeballing a number, and dragging a .pt file onto a server, then you don't have a pipeline. You have a ritual made of hopes and prayers, and it will betray you at the worst possible moment. Let's replace it with something that actually catches problems before your users do.
Ordinary software has one big bucket of tests. ML code needs at least three, and they answer three different questions. Unit tests ask "does the model code itself behave?" -- correct shapes, gradients that flow, determinism under a fixed seed. Data tests ask "is the data what we think it is?" -- schema, null rates, distributions, and the big one, leakage. Integration tests ask "does the whole pipeline run end to end without falling over?" Skip any of the three and you've left a door open. Let's start with the code itself.
# tests/test_model.py
import pytest
import torch
import numpy as np
from model import Classifier
from data import load_features, validate_schema
class TestModelArchitecture:
"""Unit tests: does the model code work correctly?"""
def test_output_shape(self):
model = Classifier(input_dim=784, n_classes=10)
x = torch.randn(32, 784)
output = model(x)
assert output.shape == (32, 10)
def test_output_probabilities(self):
model = Classifier(input_dim=784, n_classes=10)
x = torch.randn(1, 784)
probs = torch.softmax(model(x), dim=1)
assert torch.allclose(probs.sum(), torch.tensor(1.0), atol=1e-5)
def test_gradient_flow(self):
"""Ensure gradients flow through all parameters."""
model = Classifier(input_dim=784, n_classes=10)
x = torch.randn(4, 784)
y = torch.randint(0, 10, (4,))
loss = torch.nn.functional.cross_entropy(model(x), y)
loss.backward()
for name, param in model.named_parameters():
assert param.grad is not None, f"No gradient for {name}"
assert not torch.all(param.grad == 0), f"Zero gradient for {name}"
def test_deterministic_with_seed(self):
torch.manual_seed(42)
model1 = Classifier(input_dim=784, n_classes=10)
out1 = model1(torch.randn(1, 784))
torch.manual_seed(42)
model2 = Classifier(input_dim=784, n_classes=10)
out2 = model2(torch.randn(1, 784))
assert torch.allclose(out1, out2)
Look at what these tiny tests actually buy you. The shape test catches the single most common bug in deep learning -- a transposed dimension or an off-by-one in the output layer -- in milliseconds, before you waste an hour of GPU time discovering it. The gradient-flow test is my personal favourite: it catches the silent killer where a detach() snuck in somewhere, or a layer got frozen by accident, and half your network is quietly not learning a thing. The training loss goes down (the other half is compensating), everything looks fine, and your model is secretly running at half capacity. A three-line test finds it instantly. And the determinism test is the guardian of reproducibility we fought so hard for in #119 -- if the same seed gives you different outputs, something in your stack is non-deterministic and your "reproducible" experiments are a fiction.
Now the data tests, which are where ML testing really earns its keep:
# tests/test_data.py
import numpy as np
from data import load_features
class TestDataValidation:
"""Data tests: is the data what we expect?"""
def test_schema(self):
df = load_features("data/train.parquet")
expected_columns = ['user_id', 'feature_1', 'feature_2', 'label']
assert all(col in df.columns for col in expected_columns)
def test_no_nulls_in_critical_features(self):
df = load_features("data/train.parquet")
critical = ['feature_1', 'feature_2', 'label']
for col in critical:
null_rate = df[col].isnull().mean()
assert null_rate < 0.01, f"{col} has {null_rate:.1%} nulls"
def test_label_distribution(self):
df = load_features("data/train.parquet")
class_counts = df['label'].value_counts(normalize=True)
# No class should be less than 1% -- extreme imbalance
assert class_counts.min() > 0.01
def test_feature_ranges(self):
df = load_features("data/train.parquet")
assert df['feature_1'].between(-10, 10).all()
assert df['feature_2'].between(0, 1).all()
def test_no_data_leakage(self):
"""Train and test sets must not share entities."""
train = load_features("data/train.parquet")
test = load_features("data/test.parquet")
train_ids = set(train['user_id'])
test_ids = set(test['user_id'])
overlap = train_ids & test_ids
assert len(overlap) == 0, f"{len(overlap)} users appear in both sets"
That last test -- test_no_data_leakage -- is the single most important test that most ML projects don't have, and I want you to feel a little bit angry about that. Way back in episode #14 we talked about how leakage silently inflates your metrics: the same user (or the same near-duplicate row) ends up in both train and test, the model effectively memorises the answer, and your evaluation reports a gorgeous 97% that evaporates the moment it meets a real stranger in production. The nasty part is that leakage usually creeps in later, when someone refreshes the dataset and the splitting logic quietly stops holding. An automated test that runs on every data refresh catches it every single time, forever. Having said that, notice these data tests don't care about your model at all -- they're testing the world your model has to live in, and they'd catch a broken upstream pipeline (a renamed column, a unit change from [0,1] to [0,100]) before it ever poisoned a training run.
The third kind, integration tests, is the one people skip because it feels like overkill -- right up until the day the pieces work individually but explode when bolted together. An integration test runs the whole tiny pipeline end to end on a handful of rows: load, train a couple of steps, save, reload, predict. It doesn't care whether the model is any good (that's the validation gate's job); it only asserts that the machinery doesn't fall apart.
# tests/test_integration.py
import torch
from pipeline import train_step, save_model, load_model
from model import Classifier
def test_train_save_load_roundtrip(tmp_path):
"""The whole pipeline must survive a save/load without changing its mind."""
model = Classifier(input_dim=784, n_classes=10)
x = torch.randn(8, 784)
# One training step must not raise and must return a finite loss
loss = train_step(model, x, torch.randint(0, 10, (8,)))
assert torch.isfinite(torch.tensor(loss)), "training produced a non-finite loss"
# Save, reload, and confirm the reloaded model predicts identically
path = tmp_path / "model.pt"
save_model(model, path)
reloaded = load_model(Classifier(input_dim=784, n_classes=10), path)
with torch.no_grad():
before = model(x)
after = reloaded(x)
assert torch.allclose(before, after, atol=1e-6), "save/load changed the outputs"
That save-and-reload roundtrip catches a genuinely common production bug: the model that scores beautifully in memory but predicts nonsense once it's been serialised and reloaded on the serving box, because a preprocessing step or a buffer didn't get persisted. Cheap test, expensive bug ;-)
Here's a mistake I've watched quit some teams make: a training job finishes without errors, so they ship the model. But "training didn't crash" and "this model is good enough for users" are wildly different claims. A model can train perfectly and still be a disaster -- overfit, biased, slower than the SLA, or simply worse than the one already in production. So between "training succeeded" and "deploy", you put a gate. Nothing passes the gate unless it earns it.
class ModelValidator:
"""Gate model deployment on quality criteria."""
def __init__(self, thresholds):
self.thresholds = thresholds
def validate(self, metrics, production_metrics=None):
"""Check if a candidate model meets deployment criteria."""
results = []
# Absolute thresholds -- hard floors the model must clear
for metric, threshold in self.thresholds.items():
if metric in metrics:
# latency is "lower is better", everything else "higher is better"
if 'latency' in metric:
passed = metrics[metric] <= threshold
else:
passed = metrics[metric] >= threshold
results.append({
'check': f"{metric} vs {threshold}",
'value': metrics[metric],
'passed': passed,
})
# Relative check: never ship a model worse than the one already live
if production_metrics:
for metric in ['accuracy', 'f1', 'auc']:
if metric in metrics and metric in production_metrics:
diff = metrics[metric] - production_metrics[metric]
passed = diff > -0.01 # tolerate at most a 1% regression
results.append({
'check': f"{metric} not worse than production",
'new': metrics[metric],
'production': production_metrics[metric],
'diff': diff,
'passed': passed,
})
all_passed = all(r['passed'] for r in results)
return {'passed': all_passed, 'checks': results}
# Define the gates once, apply them forever
validator = ModelValidator(thresholds={
'accuracy': 0.90,
'precision': 0.85,
'recall': 0.80,
'f1': 0.85,
'latency_p99_ms': 100, # must answer in under 100ms
})
new_metrics = {'accuracy': 0.92, 'precision': 0.88, 'recall': 0.83,
'f1': 0.855, 'latency_p99_ms': 78}
prod_metrics = {'accuracy': 0.91, 'precision': 0.87, 'recall': 0.82,
'f1': 0.845}
result = validator.validate(new_metrics, prod_metrics)
if result['passed']:
print("Model APPROVED for deployment")
else:
print("Model BLOCKED from deployment:")
for check in result['checks']:
if not check['passed']:
print(f" FAILED: {check['check']}")
Two things make this gate genuinely useful, and both are easy to miss. First, the absolute thresholds are your non-negotiable floors -- accuracy below 90%, latency above 100ms, whatever your product actually needs (and yes, latency is a quality metric; a brilliant model that answers too slowly is a broken feature, as we hammered home in episode #121). Second, and more subtly, the relative check compares the candidate against the model currently serving traffic. This is the thing that saves you from the sneaky regression -- the new model clears every absolute bar but is nonetheless a hair worse than what you already have. Without that check, you'd happily deploy a downgrade because it was "good enough". With it, you refuse to move backwards. That 1% tolerance is a judgement call, by the way -- tighten it for a fraud model where every point matters, loosen it if the new model buys you something else (half the latency, say) that's worth a sliver of accuracy.
A model in production is not a file on someone's laptop called model_final_v2_REALLY_final.pt. (We've all seen that folder. Some of us are that folder.) It's a versioned, tracked thing that moves through stages -- development, then staging, then production, and eventually archived -- with a record of who promoted it, when, and on the strength of which metrics. That record is a model registry, and it turns "which model is live right now and why?" from a panicked Slack thread into a single lookup.
from datetime import datetime
class ModelRegistry:
"""A minimal model registry with a promotion workflow."""
VALID_TRANSITIONS = {
('development', 'staging'),
('staging', 'production'),
('production', 'archived'),
}
def __init__(self):
self.versions = {} # (name, version) -> record
def register(self, model_name, version, model_path, metrics):
"""Record a freshly trained model version in 'development'."""
record = {
'model_name': model_name,
'version': version,
'model_path': model_path,
'metrics': metrics,
'stage': 'development',
'registered_at': datetime.now().isoformat(),
'history': [],
}
self.versions[(model_name, version)] = record
return record
def promote(self, model_name, version, target_stage, validator=None):
"""Move a model to a new stage -- only if the transition is legal."""
record = self.versions[(model_name, version)]
current = record['stage']
if (current, target_stage) not in self.VALID_TRANSITIONS:
raise ValueError(f"Illegal transition: {current} -> {target_stage}")
# Re-run the quality gate before promoting to production
if target_stage == 'production' and validator is not None:
gate = validator.validate(record['metrics'])
if not gate['passed']:
raise ValueError("Model failed validation gate; promotion refused")
record['history'].append((current, target_stage, datetime.now().isoformat()))
record['stage'] = target_stage
return record
# In real life you'd lean on MLflow rather than rolling your own:
# import mlflow
# mlflow.register_model("runs://model", "classifier")
# client = mlflow.tracking.MlflowClient()
# client.transition_model_version_stage("classifier", version=3, stage="Production")
Notice two design choices I made deliberately. The registry refuses illegal transitions -- you cannot leap a model straight from development to production and skip staging, because that path is not in the VALID_TRANSITIONS set. And promoting to production re-runs the validation gate, so even if a model snuck into staging on a good day, it has to prove itself again before it touches real users. In production you'd almost certainly use MLflow's model registry (or SageMaker, or Vertex, take your pick) rather than this toy, but the concept is identical, and I always think it's healthier to see the 30 lines that make the idea work before you hide it behind a managed service ;-)
At the end of episode #123, on monitoring, I left you with a loose thread that's been bugging me ever since: we built beautiful drift detectors that scream when the world moves, and then... a human has to notice the scream, log in, kick off a training run, squint at the metrics, and push the result. That doesn't scale, and worse, it depends on a person being awake and paying attention. Retraining should be a pipeline that fires on its own -- on a schedule, on detected drift, or on a measured performance drop.
from datetime import datetime
class RetrainPipeline:
"""Automated retraining, triggered by schedule, drift, or degradation."""
def __init__(self, model_config, data_config, validator):
self.model_config = model_config
self.data_config = data_config
self.validator = validator
def should_retrain(self, drift_report, last_retrain_date):
"""Decide whether a retrain is warranted -- and say WHY."""
# 1. Time-based: never let a model get staler than a week
days_since = (datetime.now() - last_retrain_date).days
if days_since >= 7:
return True, "Scheduled retrain (7-day cycle)"
# 2. Drift-based: the input world moved (see episode #123)
if drift_report.get('critical_features_drifted', 0) > 2:
return True, "Data drift in multiple critical features"
# 3. Performance-based: accuracy actually fell
if drift_report.get('accuracy_drop', 0) > 0.05:
return True, "Performance degradation exceeded 5%"
return False, "No retrain needed"
def run(self):
"""Execute the retraining pipeline, step by step, stopping on failure."""
steps = [
("fetch_data", self.fetch_fresh_data),
("validate_data", self.validate_data),
("train", self.train_model),
("evaluate", self.evaluate_model),
("validate_model", self.validate_model),
("register", self.register_model),
]
for step_name, step_fn in steps:
print(f"Running: {step_name}")
result = step_fn()
if not result.get('success', True):
print(f"Pipeline FAILED at {step_name}: {result.get('error')}")
return False
return True
def fetch_fresh_data(self):
return {'success': True} # pull latest from the feature store
def validate_data(self):
return {'success': True} # run the data tests from earlier
def train_model(self):
return {'success': True} # train, with the run tracked (episode #119)
def evaluate_model(self):
return {'success': True} # score on a held-out set
def validate_model(self):
return {'success': True} # run the validation gate
def register_model(self):
return {'success': True} # register as 'staging', await promotion
The bit I want you to internalise is the should_retrain logic, because the naive version -- "retrain every night, unconditionally" -- is a trap. Blind nightly retraining burns compute you don't need to spend, and (this is the part people forget) it can actually hurt you, by chasing noise and churning a perfectly good model for a marginally different one every single day. Smart retraining is conditional: a scheduled floor so nothing rots past a week, plus drift and degradation triggers so you react fast when reality genuinely shifts. And notice the pipeline stops at the first failing step -- if the fresh data fails validation, you never even start training on garbage. That's not a nice-to-have. That's the whole point of building the thing.
None of the above matters if the environment your pipeline runs in drifts as much as your data does. An ML result that can't be reproduced isn't a result, it's an anecdote. The cheapest, highest-leverage discipline in the entire field is boringly simple: pin your dependencies and containerise your runtime, so the exact same code, on the exact same data, with the exact same library versions, gives you the exact same model. Every time. On your laptop, on the CI runner, and on the training box.
# Dockerfile -- one environment, everywhere: laptop, CI, and prod
FROM python:3.11-slim
# Pinned system deps first, so this layer caches
RUN apt-get update && apt-get install -y --no-install-recommends \
build-essential git && \
rm -rf /var/lib/apt/lists/*
WORKDIR /app
# Copy the LOCKED requirements and install exact versions
# requirements.lock is generated with `pip freeze` (or pip-tools / uv)
# and pins EVERY transitive dependency to an exact version.
COPY requirements.lock .
RUN pip install --no-cache-dir -r requirements.lock
COPY . .
# A fixed entrypoint so CI and humans run training the same way
ENTRYPOINT ["python", "src/train.py"]
The key line is COPY requirements.lock -- a locked file where every package, including the transitive dependencies you never think about, is pinned to an exact version (torch==2.3.1, not torch>=2.3). The difference is not academic. A loose >= means a fresh pip install three weeks from now silently pulls a newer NumPy, which changes a default, which shifts a rounding behaviour, which nudges your model weights, and suddenly you cannot reproduce last month's result and nobody knows why. Pin everything, rebuild the image only when you choose to, and your pipeline stops being haunted. This is infrastructure as code for ML: the environment isn't a thing you set up by hand and hope stays put, it's a file in your repo that you review like any other code.
Right, we've got tests, gates, a registry, conditional retraining, and a reproducible environment. Now let's wire them into an actual pipeline that a git push triggers. Here's a practical GitHub Actions workflow -- and I promise it's less scary than it looks:
# .github/workflows/ml-pipeline.yml
name: ML Pipeline
on:
push:
branches: [main]
paths:
- 'src/**'
- 'config/**'
- 'tests/**'
schedule:
- cron: '0 2 * * 1' # weekly retrain, Monday 02:00 UTC
jobs:
test:
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v4
- uses: actions/setup-python@v5
with:
python-version: '3.11'
- run: pip install -r requirements.lock
- run: pytest tests/test_model.py -v
- run: pytest tests/test_data.py -v
train:
needs: test # never train if the tests are red
runs-on: ubuntu-latest # or a self-hosted GPU runner
if: github.event_name == 'schedule' || contains(github.event.head_commit.message, '[retrain]')
steps:
- uses: actions/checkout@v4
- run: pip install -r requirements.lock
- run: python src/train.py --config config/production.yaml
- run: python src/evaluate.py --model-path outputs/model.pt
- run: python src/validate.py --thresholds config/thresholds.yaml
- uses: actions/upload-artifact@v4
with:
name: trained-model
path: outputs/model.pt
deploy:
needs: train
runs-on: ubuntu-latest
if: success()
steps:
- run: echo "Register model version in the registry"
- run: echo "Deploy to staging behind a canary"
# A manual approval gate guards the final step to production
Read the needs: and if: lines, because that's where the intelligence lives. The train job needs: test, so a single red test blocks the entire pipeline -- no training happens on code that fails its unit or data checks. The train job also only fires on the weekly schedule or when a commit message contains the [retrain] tag, which gives you a lovely on-demand escape hatch: push a normal commit and only tests run, push one that says [retrain] and you kick off a full training run without waiting for Monday. And deploy needs: train, so a model that flunks the validation gate inside train never reaches the deploy stage at all. The gate isn't a suggestion -- it's a wall the pipeline physically cannot walk through. Nota bene: I've left a manual approval step guarding production on purpose. Fully automating the road to staging is smart; fully automating the last hop to production, with real users on the other side, is a decision I want a human to sign off on for a good while yet.
needs: and if: so red tests block training and failed gates block deploys, with a manual approval guarding the last hop to production.Three to get your hands dirty with before the next one. As always, I'll walk through full solutions next time.
Write the gradient-flow guard for real. Build a tiny two-layer PyTorch Classifier, then deliberately sabotage it -- freeze one layer with requires_grad = False, or slip a stray .detach() into the forward pass. Now write the test_gradient_flow test from this episode and confirm it fails on the broken model and passes once you undo the sabotage. Write one sentence on why this test would have caught a bug that a shrinking training loss happily hid from you.
Make the validation gate refuse a regression. Take the ModelValidator and feed it a candidate that clears every absolute threshold (accuracy 0.91, precision 0.86, and so on) but is 2% worse on accuracy than the current production model. Confirm the gate blocks it on the relative check. Then adjust the tolerance so the same candidate would pass, and write a short paragraph on when you'd actually want to loosen it (hint: what if the candidate is half the latency?).
Design a conditional retrain trigger. Extend should_retrain with a fourth rule of your own -- say, "retrain if the fraction of predictions in the low-confidence band has risen above 20%". Pick a defensible threshold, wire it in, and write a paragraph justifying it. Bonus: describe one realistic situation where your trigger would fire but you'd actually be wrong to retrain (think seasonality, think a one-off event).
We now have a pipeline that tests a model, gates it, registers it, retrains it on its own, and ships it -- all reproducibly. That's a genuinely grown-up ML system. But I've been quietly cheating this whole episode, and it's time to admit it: every one of those train steps just... runs. I've said nothing about where the heavy lifting actually happens, or why a training run that takes eight hours on your laptop takes eight minutes on the right hardware. The whole reason modern AI is even possible lives one layer beneath everything we've built so far -- in the silicon that does thousands of multiplications at once -- and that's the machinery we finally crack open next time.
@scipio