CI/CD for ML Pipelines: Automating the Path to Production

Traditional software CI/CD is relatively straightforward: run tests, check types, build, deploy. ML CI/CD adds three new dimensions:

1. Data validation - is the training data still what we expect?
2. Model evaluation - is the new model actually better?
3. Behavioral testing - does the model pass sanity checks beyond aggregate metrics?

Building a good ML CI/CD pipeline means automating all three.

The ML CI/CD Pipeline Structure

Code push to main
        │
        ▼
┌──────────────────────────────────────────────────────────┐
│  Stage 1: Code Quality                                   │
│  - lint (ruff) + type check (mypy) + unit tests          │
│  - Fast: < 5 minutes                                     │
└──────────────────────────────────────────────────────────┘
        │ pass
        ▼
┌──────────────────────────────────────────────────────────┐
│  Stage 2: Data Validation                                │
│  - Schema checks on training data                        │
│  - Distribution checks (no sudden shifts)                │
│  - Fast: < 10 minutes                                    │
└──────────────────────────────────────────────────────────┘
        │ pass
        ▼
┌──────────────────────────────────────────────────────────┐
│  Stage 3: Model Training + Evaluation                    │
│  - Train on validated data                               │
│  - Evaluate against performance thresholds               │
│  - Compare to current production model                   │
│  - Slow: 10 minutes to hours depending on model          │
└──────────────────────────────────────────────────────────┘
        │ pass + beats production
        ▼
┌──────────────────────────────────────────────────────────┐
│  Stage 4: Behavioral / Slice Testing                     │
│  - Invariance tests                                      │
│  - Directional tests                                     │
│  - Slice performance tests                               │
│  - Fast: < 10 minutes                                    │
└──────────────────────────────────────────────────────────┘
        │ pass
        ▼
┌──────────────────────────────────────────────────────────┐
│  Stage 5: Deploy to Staging + Smoke Test                 │
│  - Container build + push                                │
│  - Deploy to staging environment                         │
│  - Integration test against staging                      │
└──────────────────────────────────────────────────────────┘
        │ pass
        ▼
   Deploy to Production (manual approval gate or automatic)

GitHub Actions Implementation

yamlCopy
# .github/workflows/ml-pipeline.yml
name: ML CI/CD Pipeline

on:
  push:
    branches: [main]
  pull_request:
    branches: [main]

jobs:
  code-quality:
    runs-on: ubuntu-latest
    steps:
      - uses: actions/checkout@v4
      - uses: actions/setup-python@v4
        with:
          python-version: '3.11'
          cache: 'pip'
      - run: pip install -r requirements-dev.txt
      - run: ruff check src/
      - run: mypy src/ --ignore-missing-imports
      - run: pytest tests/unit/ -v --tb=short

  data-validation:
    needs: code-quality
    runs-on: ubuntu-latest
    steps:
      - uses: actions/checkout@v4
      - uses: actions/setup-python@v4
        with:
          python-version: '3.11'
          cache: 'pip'
      - run: pip install -r requirements.txt
      - name: Validate training data
        env:
          AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }}
          AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
        run: python scripts/validate_data.py

  train-and-evaluate:
    needs: data-validation
    runs-on: ubuntu-latest
    steps:
      - uses: actions/checkout@v4
      - uses: actions/setup-python@v4
        with:
          python-version: '3.11'
          cache: 'pip'
      - run: pip install -r requirements.txt
      - name: Train model
        run: python scripts/train.py --output-dir /tmp/model_artifacts
      - name: Evaluate against thresholds
        run: python scripts/evaluate.py --model-dir /tmp/model_artifacts --fail-below-auc 0.80
      - name: Compare to production
        run: python scripts/compare_to_production.py --challenger-dir /tmp/model_artifacts
      - name: Upload artifacts
        uses: actions/upload-artifact@v3
        with:
          name: model-artifacts
          path: /tmp/model_artifacts/

  behavioral-tests:
    needs: train-and-evaluate
    runs-on: ubuntu-latest
    steps:
      - uses: actions/checkout@v4
      - uses: actions/download-artifact@v3
        with:
          name: model-artifacts
          path: /tmp/model_artifacts/
      - run: pip install -r requirements.txt
      - run: pytest tests/behavioral/ -v --model-dir /tmp/model_artifacts

  deploy-staging:
    needs: behavioral-tests
    runs-on: ubuntu-latest
    if: github.ref == 'refs/heads/main'
    steps:
      - uses: actions/checkout@v4
      - uses: actions/download-artifact@v3
        with:
          name: model-artifacts
          path: model_artifacts/
      - name: Build and push Docker image
        run: |
          docker build -t ${{ env.REGISTRY }}/churn-predictor:${{ github.sha }} .
          docker push ${{ env.REGISTRY }}/churn-predictor:${{ github.sha }}
      - name: Deploy to staging
        run: ./scripts/deploy.sh staging ${{ github.sha }}
      - name: Smoke test staging
        run: pytest tests/integration/ --base-url https://staging.api.example.com

Data Validation Script

pythonCopy
# scripts/validate_data.py
import pandas as pd
import numpy as np
import sys
import json

def validate_training_data(data_path: str, schema_path: str) -> bool:
    df = pd.read_parquet(data_path)
    with open(schema_path) as f:
        schema = json.load(f)

    errors = []

    # Schema validation
    for col, expected_type in schema['columns'].items():
        if col not in df.columns:
            errors.append(f"Missing column: {col}")
        elif str(df[col].dtype) != expected_type:
            errors.append(f"Column {col}: expected {expected_type}, got {df[col].dtype}")

    # Volume check
    min_rows = schema.get('min_rows', 1000)
    if len(df) < min_rows:
        errors.append(f"Too few rows: {len(df)} < {min_rows}")

    # Distribution checks (compare to baseline statistics)
    with open('data/baseline_stats.json') as f:
        baseline = json.load(f)

    for col in schema.get('monitored_columns', []):
        current_mean = df[col].mean()
        baseline_mean = baseline[col]['mean']
        baseline_std = baseline[col]['std']
        z_score = abs(current_mean - baseline_mean) / (baseline_std + 1e-10)
        if z_score > 3:
            errors.append(f"Severe distribution shift in {col}: z-score={z_score:.2f}")

    if errors:
        print("DATA VALIDATION FAILED:")
        for e in errors:
            print(f"  - {e}")
        return False

    print(f"Data validation passed: {len(df)} rows, {len(df.columns)} columns")
    return True

if __name__ == "__main__":
    success = validate_training_data(
        "s3://ml-data/training/latest.parquet",
        "data/schema.json"
    )
    sys.exit(0 if success else 1)

Behavioral Testing

Aggregate AUC does not catch all failure modes. Behavioral tests check specific behaviors:

pythonCopy
# tests/behavioral/test_model_behavior.py
import pytest
import joblib
import numpy as np

@pytest.fixture
def model(request):
    model_dir = request.config.getoption("--model-dir")
    return joblib.load(f"{model_dir}/pipeline.joblib")

def make_customer(age=30, income=60000, tenure_months=12):
    return np.array([[age, income, tenure_months]])

class TestInvariance:
    """Small, irrelevant input changes should not change prediction much."""
    def test_income_noise_stability(self, model):
        base = make_customer(income=60000)
        noisy = make_customer(income=60001)  # $1 difference
        base_pred = model.predict_proba(base)[0][1]
        noisy_pred = model.predict_proba(noisy)[0][1]
        assert abs(base_pred - noisy_pred) < 0.01

class TestDirectional:
    """Changing a feature in a known direction should move prediction accordingly."""
    def test_longer_tenure_reduces_churn(self, model):
        short_tenure = make_customer(tenure_months=1)
        long_tenure = make_customer(tenure_months=60)
        assert (model.predict_proba(long_tenure)[0][1] <
                model.predict_proba(short_tenure)[0][1])

    def test_higher_income_reduces_churn(self, model):
        low_income = make_customer(income=20000)
        high_income = make_customer(income=200000)
        assert (model.predict_proba(high_income)[0][1] <
                model.predict_proba(low_income)[0][1])

class TestEdgeCases:
    def test_handles_zero_tenure(self, model):
        pred = model.predict_proba(make_customer(tenure_months=0))
        assert 0 <= pred[0][1] <= 1  # valid probability

    def test_handles_extreme_income(self, model):
        pred = model.predict_proba(make_customer(income=10_000_000))
        assert 0 <= pred[0][1] <= 1

These tests encode business knowledge. A model that predicts higher churn for higher income customers is almost certainly wrong, even if its AUC looks fine.

The Production Comparison Gate

The most important check: does the new model actually beat production?

pythonCopy
# scripts/compare_to_production.py
import mlflow
import joblib
import sys
import numpy as np
from sklearn.metrics import roc_auc_score

def compare_to_production(challenger_dir: str, holdout_data: str) -> bool:
    # Load challenger (newly trained model)
    challenger = joblib.load(f"{challenger_dir}/pipeline.joblib")

    # Load champion (current production model)
    client = mlflow.tracking.MlflowClient()
    champion_uri = "models:/churn-predictor/Production"
    champion = mlflow.sklearn.load_model(champion_uri)

    # Evaluate both on holdout set
    X_hold, y_hold = load_holdout_data(holdout_data)
    challenger_auc = roc_auc_score(y_hold, challenger.predict_proba(X_hold)[:, 1])
    champion_auc = roc_auc_score(y_hold, champion.predict_proba(X_hold)[:, 1])

    print(f"Champion AUC: {champion_auc:.4f}")
    print(f"Challenger AUC: {challenger_auc:.4f}")
    print(f"Delta: {challenger_auc - champion_auc:+.4f}")

    # Require challenger to beat champion by at least 0.5%
    if challenger_auc >= champion_auc + 0.005:
        print("RESULT: Challenger wins  -  promoting to staging")
        return True
    else:
        print("RESULT: Challenger did not beat champion  -  blocking deployment")
        return False

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("--challenger-dir", required=True)
    args = parser.parse_args()
    success = compare_to_production(args.challenger_dir, "data/holdout.parquet")
    sys.exit(0 if success else 1)

This gate means you only deploy when a model is provably better, not just "recently retrained."

What to Practice Next

• Add a GitHub Actions workflow to an existing ML repo that runs pytest on your data validation and model evaluation code on every pull request - start with unit tests, then add integration tests.
• Implement a simple model quality gate: after training, assert that validation accuracy exceeds a threshold before the workflow proceeds to the push/deploy step; treat a regression as a build failure.
• Compare DVC and MLflow Model Registry for artifact versioning in a small experiment - document which abstraction layer each operates at and when you would reach for one over the other.