AI Hints
Home›MLOps Tutorial›Registry, Testing, Formats
MLOps · Lesson 14 of 80

Testing ML Systems

Tests for data, code, and models.

⏱ 2 min read▶ Video at 32:45🎓 Free
Watch this lesson in the MLOps courseWatch on YouTube ↗Full MLOps course

Software engineers test their code before shipping. Machine learning needs the same discipline, and more. You must test the code, the data, and the model itself. A good test suite is the guard at the gate of your registry. In this lesson, we write real tests with pytest.

Testing ML Systems
From the video: Testing ML Systems
  • Test the code, the data, the model
  • Automatic, on every change
  • A guard at the registry gate

Three kinds of ML tests

There are three kinds of tests. Data tests check the inputs, no missing values, sensible ranges, a plausible share of churners. Model tests check quality, is the model better than a trivial baseline, and better than the current champion? And behaviour tests check common sense. If a customer calls support more often, their churn risk should not go down. These tests catch problems that a single accuracy number hides.

Data testsmissing values, ranges, balance
Quality testsbeats the baseline
Behaviour testscommon-sense rules
Code testsordinary unit tests

Model tests

Here are model tests. The first compares our model with a baseline, that always predicts the most common answer. We use the F one score, because only about one customer in eight leaves, so a lazy model that says nobody leaves, already scores high accuracy. The model must beat the baseline clearly. The second test is a behaviour test. Two identical customers, except that one called support six times. The riskier one must receive a higher churn probability. And the third checks that outputs are valid probabilities.

tests/test_model.pypython
import pandas as pd
from sklearn.dummy import DummyClassifier
from sklearn.metrics import f1_score
from churn import make_data

def test_beats_the_baseline(model):
    test = make_data(1000, seed=2)
    X, y = test.drop(columns="churn"), test["churn"]
    baseline = DummyClassifier(strategy="most_frequent").fit(X, y)
    assert f1_score(y, model.predict(X)) > f1_score(y, baseline.predict(X)) + 0.2

def test_more_support_calls_never_lowers_risk(model):
    low = pd.DataFrame([{"tenure": 24, "charges": 70.0, "calls": 0, "contract": 0}])
    high = pd.DataFrame([{"tenure": 24, "charges": 70.0, "calls": 6, "contract": 0}])
    assert model.predict_proba(high)[0][1] > model.predict_proba(low)[0][1]

def test_prediction_is_a_probability(model):
    p = model.predict_proba(pd.DataFrame([{"tenure": 12, "charges": 50.0, "calls": 2, "contract": 1}]))[0][1]
    assert 0.0 <= p <= 1.0

Running the tests

Let's run the suite. Pytest finds the tests, runs them, and prints a dot for every pass. All six pass. Now let's inject bad data on purpose, a negative monthly charge. We run the data tests again. This time, a test fails, and pytest shows exactly which check broke, and the offending value. In a real pipeline, this failure would stop the process, before a bad model ever reached the registry.

terminal
pytest -q -p no:cacheprovider -W ignore
......                                                                   [100%]
6 passed in 0.05s
INJECT_BAD_DATA=1 pytest -q -p no:cacheprovider -W ignore tests/test_data.py 2>&1 | grep -E 'assert|passed|failed' | head -4
>       assert (data["charges"] > 0).all()
E       assert np.False_
FAILED tests/test_data.py::test_charges_are_positive - assert np.False_
1 failed, 2 passed in 0.08s
File used: conftest.py
import os
import pytest
from sklearn.linear_model import LogisticRegression
from churn import make_data

@pytest.fixture(scope="session")
def data():
    df = make_data(3000, seed=1)
    if os.environ.get("INJECT_BAD_DATA"):
        df.loc[0, "charges"] = -5.0
    return df

@pytest.fixture(scope="session")
def model(data):
    X, y = data.drop(columns="churn"), data["churn"]
    return LogisticRegression(max_iter=500).fit(X, y)
File used: churn.py
import numpy as np, pandas as pd

def make_data(n, seed, shifted=False):
    rng = np.random.default_rng(seed)
    tenure = rng.integers(1, 72, n)
    charges = rng.normal(65 + (28 if shifted else 0), 18, n).clip(20, 140).round(2)
    calls = rng.poisson(1.5 + (1.5 if shifted else 0), n)
    contract = rng.choice([0, 1, 2], n, p=[.55, .25, .20] if not shifted else [.30, .30, .40])
    logit = -1.0 - 0.04 * tenure + 0.025 * (charges - 65) + 0.45 * calls - 0.8 * contract
    if shifted:
        logit = -0.2 - 0.01 * tenure - 0.01 * (charges - 65) + 0.2 * calls - 0.2 * contract
    churn = (rng.random(n) < 1 / (1 + np.exp(-logit))).astype(int)
    return pd.DataFrame({"tenure": tenure, "charges": charges, "calls": calls, "contract": contract, "churn": churn})
File used: tests/test_data.py
def test_no_missing_values(data):
    assert not data.isna().any().any()

def test_charges_are_positive(data):
    assert (data["charges"] > 0).all()

def test_churn_rate_is_plausible(data):
    assert 0.1 < data["churn"].mean() < 0.6
File used: tests/test_model.py
import pandas as pd
from sklearn.dummy import DummyClassifier
from sklearn.metrics import f1_score
from churn import make_data

def test_beats_the_baseline(model):
    test = make_data(1000, seed=2)
    X, y = test.drop(columns="churn"), test["churn"]
    baseline = DummyClassifier(strategy="most_frequent").fit(X, y)
    assert f1_score(y, model.predict(X)) > f1_score(y, baseline.predict(X)) + 0.2

def test_more_support_calls_never_lowers_risk(model):
    low = pd.DataFrame([{"tenure": 24, "charges": 70.0, "calls": 0, "contract": 0}])
    high = pd.DataFrame([{"tenure": 24, "charges": 70.0, "calls": 6, "contract": 0}])
    assert model.predict_proba(high)[0][1] > model.predict_proba(low)[0][1]

def test_prediction_is_a_probability(model):
    p = model.predict_proba(pd.DataFrame([{"tenure": 12, "charges": 50.0, "calls": 2, "contract": 1}]))[0][1]
    assert 0.0 <= p <= 1.0

Accuracy can lie

One more lesson from that baseline test. Accuracy can lie. When most customers stay, a model that always says stay, is right most of the time, and totally useless. For imbalanced problems, use precision, recall, and F one, or the area under the curve. Always compare against a trivial baseline, and always ask who is being missed. A test that includes a baseline, is a test that cannot be fooled by a flattering number.

Accuracy alone

  • Hides class imbalance
  • A lazy model scores high
  • Looks great, helps nobody

Better metrics

  • Precision, recall, F1
  • Compare against a baseline
  • Look at who is missed

Summary

Excellent. Test the data, the quality of the model, and its behaviour. Beware of accuracy on imbalanced problems, and always compare with a baseline. These tests guard the registry gate. Next, getting the model out of Python, into portable formats.

Key takeaways
  • Data, quality, behaviour tests
  • Compare with a baseline
  • Use F1 on imbalanced data
Next up: Model Formats — pickle, ONNX, and portability