The Testing Problem Nobody Talks About
CI/CD for machine learning is harder than CI/CD for software, and it's not because the code is more complex. It's because you're testing two things at once — whether the code runs correctly, and whether the model it produces is any good. Traditional software testing gives you the first part. The second part requires entirely different infrastructure.
Most ML teams I've worked with start by adding model training to their existing CI pipeline. They run pytest, train a small model, check that metrics exceed a threshold, and call it done. This works for about three months before someone discovers that the CI model and the production model have nothing in common because the CI training data is a 1% sample that doesn't represent the actual distribution.
Test Pyramid for ML Code
The traditional test pyramid — lots of unit tests, fewer integration tests, a handful of end-to-end tests — still applies to ML code. But the layers test different things than in conventional software.
# Unit test: feature transformation logic
def test_normalize_transaction_amount():
raw = pd.DataFrame({"amount": [10.0, 100.0, 1000.0]})
result = normalize_amount(raw)
assert result["amount_log"].iloc[0] == pytest.approx(2.302, abs=0.01)
assert result["amount_zscore"].mean() == pytest.approx(0.0, abs=0.01)
# Unit test: model input validation
def test_feature_schema_enforcement():
schema = load_feature_schema("v3.2")
sample = generate_sample_features()
sample = sample.drop(columns=["merchant_category"])
with pytest.raises(SchemaValidationError, match="merchant_category"):
schema.validate(sample)
# Integration test: training pipeline produces valid artifact
def test_training_pipeline_produces_loadable_model(tmp_path):
config = TrainingConfig(
data_path="tests/fixtures/sample_100k.parquet",
output_dir=str(tmp_path),
epochs=2,
batch_size=64,
)
train(config)
model = load_model(tmp_path / "model.pt")
assert model isn't None
sample = torch.randn(1, 147)
output = model(sample)
assert output.shape == (1, 1)Data Validation in the Pipeline
Code can be correct and the model still terrible because the data shifted. Your CI pipeline needs data validation gates, not just code validation gates.
We run three data checks before training starts. First, schema validation — column names, types, and nullable constraints match the expected schema. Second, distribution checks — feature means and standard deviations haven't drifted more than two standard deviations from baseline. Third, completeness checks — the null rate per column stays within historical bounds.
# Data validation gate in CI
from data_validator import DataValidator
validator = DataValidator(
baseline_stats="configs/feature_baselines_2024q3.json"
)
report = validator.validate("data/training_current.parquet")
if report.schema_failures:
raise PipelineError(f"Schema violations: {report.schema_failures}")
if report.drift_warnings:
logger.warning(f"Feature drift detected: {report.drift_warnings}")
notify_slack(f"Training data drift: {len(report.drift_warnings)} features")
if report.completeness_ratio < 0.95:
raise PipelineError(
f"Data completeness {report.completeness_ratio:.2%} below 95% threshold"
)That completeness check caught a data pipeline bug for us last quarter. An upstream table's ETL job started silently dropping records with certain currency codes. Our model's training data went from 2.8 million rows to 2.1 million overnight, and the completeness gate blocked the training run before we produced a model trained on biased data.
For a related perspective, see Time Series Forecasting with Temporal Fusion Transformers.
Model Quality Gates
After training completes, the model needs to pass quality gates before it's registered. These aren't the same as the metrics you track during training. Training metrics tell you how well the model fits the data. Quality gates tell you whether the model is safe to deploy.
Our quality gate checks include: overall performance above a minimum threshold (not above the champion — just above the floor), performance on critical subgroups within acceptable bounds, prediction distribution similarity to the champion model, inference latency within the serving SLA, and model size within the deployment budget.
The subgroup check is the one most teams skip and shouldn't. A model can have great aggregate AUC but perform terribly on a specific customer segment. If that segment happens to be a protected class, you've got a compliance problem. If it's a high-value customer segment, you've got a revenue problem.
Pipeline Orchestration
The CI/CD pipeline for ML has more stages than a traditional pipeline. A typical flow: lint and unit tests (2 minutes), data validation (5 minutes), training on sampled data (20 minutes for CI, hours for production), model quality gates (10 minutes), registry promotion (1 minute), staging deployment (5 minutes), integration tests against staging (10 minutes).
Running this in a single linear pipeline takes about an hour. That's too slow for development iteration but appropriate for the merge-to-main flow. Feature branches run only through data validation and a quick training pass. The full pipeline runs on merge to main and on scheduled nightly builds with full-size data.
This connects to the ideas in Object Detection Model Selection: YOLO vs DETR vs EfficientD.
# GitHub Actions workflow for ML CI/CD
name: ml-pipeline
on:
push:
branches: [main]
pull_request:
branches: [main]
jobs:
lint-and-test:
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v4
- run: pip install -r requirements-dev.txt
- run: ruff check .
- run: pytest tests/unit/ -v
data-validation:
needs: lint-and-test
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v4
- run: python scripts/validate_data.py --config configs/data_schema.yaml
train-ci:
needs: data-validation
runs-on: [self-hosted, gpu]
if: github.event_name == 'push'
steps:
- uses: actions/checkout@v4
- run: python train.py --config configs/ci_training.yaml
- run: python scripts/quality_gates.py --model outputs/model.ptRollback Strategy
Deployment isn't the end of the pipeline. You need automated rollback triggers that revert to the previous model version when production metrics degrade. We use a canary deployment that sends 5% of traffic to the new model for four hours, comparing prediction quality and latency against the incumbent. If the canary's metrics drop below the threshold, automatic rollback kicks in.
The tricky part is defining what "degraded" means in real time. Offline metrics like AUC require labels, which arrive days or weeks after prediction time. We use proxy metrics instead — prediction distribution divergence, feature importance stability, and downstream business metrics that correlate with model quality. It's not perfect, but it catches the catastrophic failures that matter most in the first few hours after deployment.
One lesson I'll pass along: always test the rollback mechanism itself. We had a rollback configuration that pointed to a model version that had been garbage-collected from the registry. The first time we needed an actual rollback, it failed, and we spent thirty minutes manually deploying the previous version while the broken model served production traffic. Now we verify rollback targets weekly as part of routine maintenance.
Testing Data Dependencies
ML pipelines have data dependencies that software pipelines don't. Your model training depends on a specific data schema, specific feature distributions, and specific data volumes. When any of these change upstream, the training pipeline might still succeed — it just produces a worse model.
We added data contract tests to the CI pipeline. These tests verify that the interface between the data pipeline and the training pipeline matches the agreed-upon contract. Column names and types must match the schema. Feature value ranges must fall within expected bounds. Row counts must exceed minimum thresholds for statistical validity.
We covered a related topic in Optical Character Recognition Pipeline Design for Noisy Docu.
# Data contract test
def test_training_data_contract():
data = load_training_data("tests/fixtures/sample_data.parquet")
# Schema contract
required_cols = ["user_id", "amount", "merchant_category", "timestamp"]
assert all(col in data.columns for col in required_cols)
# Type contract
assert data["amount"].dtype == "float64"
assert data["user_id"].dtype == "object"
# Volume contract
assert len(data) >= 50_000, f"Too few rows: {len(data)}"
# Distribution contract
assert 0 < data["amount"].mean() < 10_000
assert data["merchant_category"].nunique() >= 10Environment Reproducibility
Training environment differences cause more failed deployments than algorithm bugs. A model trained on CUDA 12.1 with PyTorch 2.1 might produce different numerical results (or crash entirely) when served on CUDA 11.8 with PyTorch 2.0. Pinning exact versions isn't just good practice — it's a correctness requirement.
We use Docker images for training and serving environments, with the Dockerfile checked into the same repository as the training code. The CI pipeline builds the Docker image as the first step, and all subsequent pipeline stages run inside that image. The image tag includes the git commit hash, so every model artifact links back to the exact environment that produced it.
For GPU-specific compatibility, we maintain a compatibility matrix that maps CUDA versions to driver versions to PyTorch versions. New combinations get tested before they're added to the matrix. We learned this lesson after a driver update on our training cluster silently changed the behavior of certain cuDNN operations, producing models that were subtly different from their predecessors. The metric difference was small enough to pass quality gates but large enough to confuse users who noticed their recommendations changing.
Branch-Based Model Development
Most ML teams develop models on feature branches and merge to main when the experiment succeeds. But the CI pipeline needs to handle both scenarios: quick validation on feature branches (is the code correct?) and full pipeline runs on main (is the model good enough to deploy?).
Feature branch pipelines run lint, unit tests, and a quick training pass on a data sample. Main branch pipelines run the full training pass with production data, quality gates, and automated registration. This split keeps feature branch iteration fast (under 15 minutes) while ensuring that main branch merges meet the full quality bar.
The risk is that a feature branch's quick training pass gives misleading signals. A model that looks great on a 10% data sample might fail on the full dataset because the sample missed a critical edge case. We mitigate this by requiring that feature branches pass a "representative sample" test — the sample must contain at least one example from each category in the stratification key, and the sample distribution must not diverge from the full dataset distribution by more than a PSI of 0.1.