Ml Pipeline Automation

ML Pipeline Automation

Orchestrate end-to-end machine learning workflows from data ingestion to production deployment with production-tested Airflow, Kubeflow, and MLflow patterns.

When to Use This Skill

Load this skill when:

• Building ML Pipelines: Orchestrating data → train → deploy workflows
• Scheduling Retraining: Setting up automated model retraining schedules
• Experiment Tracking: Tracking experiments, parameters, metrics across runs
• MLOps Implementation: Building reproducible, monitored ML infrastructure
• Workflow Orchestration: Managing complex multi-step ML workflows
• Model Registry: Managing model versions and deployment lifecycle

Quick Start: ML Pipeline in 5 Steps

# 1. Install Airflow and MLflow (check for latest versions at time of use)
pip install apache-airflow==3.1.5 mlflow==3.7.0

# Note: These versions are current as of December 2025
# Check PyPI for latest stable releases: https://pypi.org/project/apache-airflow/

# 2. Initialize Airflow database
airflow db init

# 3. Create DAG file: dags/ml_training_pipeline.py
cat > dags/ml_training_pipeline.py << 'EOF'
from airflow import DAG
from airflow.operators.python import PythonOperator
from datetime import datetime, timedelta

default_args = {
    'owner': 'ml-team',
    'retries': 2,
    'retry_delay': timedelta(minutes=5)
}

dag = DAG(
    'ml_training_pipeline',
    default_args=default_args,
    schedule_interval='@daily',
    start_date=datetime(2025, 1, 1)
)

def train_model(**context):
    import mlflow
    import mlflow.sklearn
    from sklearn.ensemble import RandomForestClassifier
    from sklearn.datasets import load_iris
    from sklearn.model_selection import train_test_split

    X, y = load_iris(return_X_y=True)
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2)

    mlflow.set_tracking_uri('http://localhost:5000')
    mlflow.set_experiment('iris-training')

    with mlflow.start_run():
        model = RandomForestClassifier(n_estimators=100)
        model.fit(X_train, y_train)

        accuracy = model.score(X_test, y_test)
        mlflow.log_metric('accuracy', accuracy)
        mlflow.sklearn.log_model(model, 'model')

train = PythonOperator(
    task_id='train_model',
    python_callable=train_model,
    dag=dag
)
EOF

# 4. Start Airflow scheduler and webserver
airflow scheduler &
airflow webserver --port 8080 &

# 5. Trigger pipeline
airflow dags trigger ml_training_pipeline

# Access UI: http://localhost:8080

Result: Working ML pipeline with experiment tracking in under 5 minutes.

Core Concepts

Pipeline Stages

1. Data Collection → Fetch raw data from sources 2. Data Validation → Check schema, quality, distributions 3. Feature Engineering → Transform raw data to features 4. Model Training → Train with hyperparameter tuning 5. Model Evaluation → Validate performance on test set 6. Model Deployment → Push to production if metrics pass 7. Monitoring → Track drift, performance in production

Orchestration Tools Comparison

Basic Airflow DAG

from airflow import DAG
from airflow.operators.python import PythonOperator
from datetime import datetime, timedelta
import logging

logger = logging.getLogger(__name__)

default_args = {
    'owner': 'ml-team',
    'depends_on_past': False,
    'email': ['alerts@example.com'],
    'email_on_failure': True,
    'retries': 2,
    'retry_delay': timedelta(minutes=5)
}

dag = DAG(
    'ml_training_pipeline',
    default_args=default_args,
    description='End-to-end ML training pipeline',
    schedule_interval='@daily',
    start_date=datetime(2025, 1, 1),
    catchup=False
)

def validate_data(**context):
    """Validate input data quality."""
    import pandas as pd

    data_path = "/data/raw/latest.csv"
    df = pd.read_csv(data_path)

    # Validation checks
    assert len(df) > 1000, f"Insufficient data: {len(df)} rows"
    assert df.isnull().sum().sum() < len(df) * 0.1, "Too many nulls"

    context['ti'].xcom_push(key='data_path', value=data_path)
    logger.info(f"Data validation passed: {len(df)} rows")

def train_model(**context):
    """Train ML model with MLflow tracking."""
    import mlflow
    import mlflow.sklearn
    from sklearn.ensemble import RandomForestClassifier

    data_path = context['ti'].xcom_pull(key='data_path', task_ids='validate_data')

    mlflow.set_tracking_uri('http://mlflow:5000')
    mlflow.set_experiment('production-training')

    with mlflow.start_run():
        # Training logic here
        model = RandomForestClassifier(n_estimators=100)
        # model.fit(X, y) ...

        mlflow.log_param('n_estimators', 100)
        mlflow.sklearn.log_model(model, 'model')

validate = PythonOperator(
    task_id='validate_data',
    python_callable=validate_data,
    dag=dag
)

train = PythonOperator(
    task_id='train_model',
    python_callable=train_model,
    dag=dag
)

validate >> train

Known Issues Prevention

1. Task Failures Without Alerts

Problem: Pipeline fails silently, no one notices until users complain.

Solution: Configure email/Slack alerts on failure:

default_args = {
    'email': ['ml-team@example.com'],
    'email_on_failure': True,
    'email_on_retry': False
}

def on_failure_callback(context):
    """Send Slack alert on failure."""
    from airflow.providers.slack.operators.slack_webhook import SlackWebhookOperator

    slack_msg = f"""
    :red_circle: Task Failed: {context['task_instance'].task_id}
    DAG: {context['task_instance'].dag_id}
    Execution Date: {context['ds']}
    Error: {context.get('exception')}
    """

    SlackWebhookOperator(
        task_id='slack_alert',
        slack_webhook_conn_id='slack_webhook',
        message=slack_msg
    ).execute(context)

task = PythonOperator(
    task_id='critical_task',
    python_callable=my_function,
    on_failure_callback=on_failure_callback,
    dag=dag
)

2. Missing XCom Data Between Tasks

Problem: Task expects XCom value from previous task, gets None, crashes.

Solution: Always validate XCom pulls:

def process_data(**context):
    data_path = context['ti'].xcom_pull(
        key='data_path',
        task_ids='upstream_task'
    )

    if data_path is None:
        raise ValueError("No data_path from upstream_task - check XCom push")

    # Process data...

3. DAG Not Appearing in UI

Problem: DAG file exists in dags/ but doesn't show in Airflow UI.

Solution: Check DAG parsing errors:

# Check for syntax errors
python dags/my_dag.py

# View DAG import errors in UI
# Navigate to: Browse → DAG Import Errors

# Common fixes:
# 1. Ensure DAG object is defined in file
# 2. Check for circular imports
# 3. Verify all dependencies installed
# 4. Fix syntax errors

4. Hardcoded Paths Break in Production

Problem: Paths like /Users/myname/data/ work locally, fail in production.

Solution: Use Airflow Variables or environment variables:

from airflow.models import Variable

def load_data(**context):
    # ❌ Bad: Hardcoded path
    # data_path = "/Users/myname/data/train.csv"

    # ✅ Good: Use Airflow Variable
    data_dir = Variable.get("data_directory", "/data")
    data_path = f"{data_dir}/train.csv"

    # Or use environment variable
    import os
    data_path = os.getenv("DATA_PATH", "/data/train.csv")

5. Stuck Tasks Consume Resources

Problem: Task hangs indefinitely, blocks worker slot, wastes resources.

Solution: Set execution_timeout on tasks:

from datetime import timedelta

task = PythonOperator(
    task_id='long_running_task',
    python_callable=my_function,
    execution_timeout=timedelta(hours=2),  # Kill after 2 hours
    dag=dag
)

6. No Data Validation = Bad Model Training

Problem: Train on corrupted/incomplete data, model performs poorly in production.

Solution: Add data quality validation tasks:

def validate_data_quality(**context):
    """Comprehensive data validation."""
    import pandas as pd

    df = pd.read_csv(data_path)

    # Schema validation
    required_cols = ['user_id', 'timestamp', 'feature_a', 'target']
    missing_cols = set(required_cols) - set(df.columns)
    if missing_cols:
        raise ValueError(f"Missing columns: {missing_cols}")

    # Statistical validation
    if df['target'].isnull().sum() > 0:
        raise ValueError("Target column contains nulls")

    if len(df) < 1000:
        raise ValueError(f"Insufficient data: {len(df)} rows")

    logger.info("✅ Data quality validation passed")

7. Untracked Experiments = Lost Knowledge

Problem: Can't reproduce results, don't know which hyperparameters worked.

Solution: Use MLflow for all experiments:

import mlflow

mlflow.set_tracking_uri('http://mlflow:5000')
mlflow.set_experiment('model-experiments')

with mlflow.start_run(run_name='rf_v1'):
    # Log ALL hyperparameters
    mlflow.log_params({
        'model_type': 'random_forest',
        'n_estimators': 100,
        'max_depth': 10,
        'random_state': 42
    })

    # Log ALL metrics
    mlflow.log_metrics({
        'train_accuracy': 0.95,
        'test_accuracy': 0.87,
        'f1_score': 0.89
    })

    # Log model
    mlflow.sklearn.log_model(model, 'model')

When to Load References

Load reference files for detailed production implementations:

• Airflow DAG Patterns: Load references/airflow-patterns.md when building complex DAGs with error handling, dynamic generation, sensors, task groups, or retry logic. Contains complete production DAG examples.

• Kubeflow & MLflow Integration: Load references/kubeflow-mlflow.md when using Kubeflow Pipelines for container-native orchestration, integrating MLflow tracking, building KFP components, or managing model registry.

• Pipeline Monitoring: Load references/pipeline-monitoring.md when implementing data quality checks, drift detection, alert configuration, or pipeline health monitoring with Prometheus.

Best Practices

1. Idempotent Tasks: Tasks should produce same result when re-run 2. Atomic Operations: Each task does one thing well 3. Version Everything: Data, code, models, dependencies 4. Comprehensive Logging: Log all important events with context 5. Error Handling: Fail fast with clear error messages 6. Monitoring: Track pipeline health, data quality, model drift 7. Testing: Test tasks independently before integrating 8. Documentation: Document DAG purpose, task dependencies

Common Patterns

Conditional Execution

from airflow.operators.python import BranchPythonOperator

def choose_branch(**context):
    accuracy = context['ti'].xcom_pull(key='accuracy', task_ids='evaluate')

    if accuracy > 0.9:
        return 'deploy_to_production'
    else:
        return 'retrain_with_more_data'

branch = BranchPythonOperator(
    task_id='check_accuracy',
    python_callable=choose_branch,
    dag=dag
)

train >> evaluate >> branch >> [deploy, retrain]

Parallel Training

from airflow.utils.task_group import TaskGroup

with TaskGroup('train_models', dag=dag) as train_group:
    train_rf = PythonOperator(task_id='train_rf', ...)
    train_lr = PythonOperator(task_id='train_lr', ...)
    train_xgb = PythonOperator(task_id='train_xgb', ...)

# All models train in parallel
preprocess >> train_group >> select_best

Waiting for Data

from airflow.sensors.filesystem import FileSensor

wait_for_data = FileSensor(
    task_id='wait_for_data',
    filepath='/data/input/{{ ds }}.csv',
    poke_interval=60,  # Check every 60 seconds
    timeout=3600,  # Timeout after 1 hour
    mode='reschedule',  # Don't block worker
    dag=dag
)

wait_for_data >> process_data

Airflow DAG Patterns for ML Pipelines

Production-ready Airflow patterns for ML workflows with error handling, retries, sensors, and dynamic DAG generation.

Complete ML Training DAG

from airflow import DAG
from airflow.operators.python import PythonOperator
from airflow.operators.bash import BashOperator
from airflow.sensors.filesystem import FileSensor
from airflow.utils.task_group import TaskGroup
from datetime import datetime, timedelta
import logging

logger = logging.getLogger(__name__)

# Default arguments with retry logic
default_args = {
    'owner': 'ml-team',
    'depends_on_past': False,
    'email': ['alerts@example.com'],
    'email_on_failure': True,
    'email_on_retry': False,
    'retries': 2,
    'retry_delay': timedelta(minutes=5),
    'retry_exponential_backoff': True,
    'max_retry_delay': timedelta(minutes=30)
}

dag = DAG(
    'ml_training_pipeline',
    default_args=default_args,
    description='Complete ML training pipeline',
    schedule_interval='0 2 * * *',  # Daily at 2 AM
    start_date=datetime(2025, 1, 1),
    catchup=False,  # Don't backfill
    max_active_runs=1,  # One run at a time
    tags=['ml', 'training']
)

# Task 1: Data Validation
def validate_data(**context):
    """Validate input data quality."""
    from datetime import date
    import pandas as pd

    execution_date = context['ds']
    data_path = f"/data/raw/{execution_date}.csv"

    try:
        df = pd.read_csv(data_path)

        # Validation checks
        assert len(df) > 1000, f"Insufficient data: {len(df)} rows"
        assert df.isnull().sum().sum() < len(df) * 0.1, "Too many nulls"

        # Push metadata to XCom
        context['ti'].xcom_push(key='row_count', value=len(df))
        context['ti'].xcom_push(key='data_path', value=data_path)

        logger.info(f"Data validation passed: {len(df)} rows")

    except Exception as e:
        logger.error(f"Data validation failed: {e}")
        raise

validate = PythonOperator(
    task_id='validate_data',
    python_callable=validate_data,
    dag=dag
)

# Task 2: Feature Engineering
def engineer_features(**context):
    """Create features from raw data."""
    import pandas as pd
    import numpy as np
    from sklearn.preprocessing import StandardScaler

    data_path = context['ti'].xcom_pull(key='data_path', task_ids='validate_data')
    df = pd.read_csv(data_path)

    # Feature engineering
    df['feature_ratio'] = df['feature_a'] / (df['feature_b'] + 1e-6)
    df['feature_log'] = np.log1p(df['feature_c'])

    # Save processed data
    processed_path = data_path.replace('raw', 'processed')
    df.to_csv(processed_path, index=False)

    context['ti'].xcom_push(key='processed_path', value=processed_path)
    logger.info(f"Feature engineering complete: {processed_path}")

features = PythonOperator(
    task_id='engineer_features',
    python_callable=engineer_features,
    dag=dag
)

# Task Group: Model Training
with TaskGroup('train_models', dag=dag) as train_group:

    def train_model(model_type, **context):
        """Train a specific model type."""
        import mlflow
        import pandas as pd
        from sklearn.ensemble import RandomForestClassifier
        from sklearn.linear_model import LogisticRegression
        import joblib

        processed_path = context['ti'].xcom_pull(
            key='processed_path',
            task_ids='engineer_features'
        )

        df = pd.read_csv(processed_path)
        X = df.drop('target', axis=1)
        y = df['target']

        # Select model
        if model_type == 'rf':
            model = RandomForestClassifier(n_estimators=100)
        elif model_type == 'lr':
            model = LogisticRegression()

        # Train with MLflow tracking
        mlflow.set_experiment('ml_pipeline')

        with mlflow.start_run(run_name=f"{model_type}_{context['ds']}"):
            model.fit(X, y)

            # Log metrics
            accuracy = model.score(X, y)
            mlflow.log_metric('accuracy', accuracy)
            mlflow.log_param('model_type', model_type)

            # Save model
            model_path = f"/models/{model_type}_{context['ds']}.pkl"
            joblib.dump(model, model_path)
            mlflow.log_artifact(model_path)

            # Save run_id to XCom for deployment
            run_id = mlflow.active_run().info.run_id
            context['ti'].xcom_push(key=f'{model_type}_run_id', value=run_id)

            context['ti'].xcom_push(key=f'{model_type}_accuracy', value=accuracy)
            logger.info(f"{model_type} trained: accuracy={accuracy:.4f}")

    train_rf = PythonOperator(
        task_id='train_random_forest',
        python_callable=lambda **ctx: train_model('rf', **ctx),
        dag=dag
    )

    train_lr = PythonOperator(
        task_id='train_logistic_regression',
        python_callable=lambda **ctx: train_model('lr', **ctx),
        dag=dag
    )

# Task: Model Selection
def select_best_model(**context):
    """Select best performing model."""
    rf_acc = context['ti'].xcom_pull(
        key='rf_accuracy',
        task_ids='train_models.train_random_forest'
    )
    lr_acc = context['ti'].xcom_pull(
        key='lr_accuracy',
        task_ids='train_models.train_logistic_regression'
    )

    best_model = 'rf' if rf_acc > lr_acc else 'lr'
    best_acc = max(rf_acc, lr_acc)

    context['ti'].xcom_push(key='best_model', value=best_model)
    context['ti'].xcom_push(key='best_accuracy', value=best_acc)

    logger.info(f"Best model: {best_model} (accuracy={best_acc:.4f})")

    # Fail if accuracy too low
    if best_acc < 0.8:
        raise ValueError(f"Best accuracy {best_acc:.4f} below threshold 0.8")

select = PythonOperator(
    task_id='select_best_model',
    python_callable=select_best_model,
    dag=dag
)

# Task: Deploy Model
def deploy_model(**context):
    """Deploy best model to production."""
    best_model = context['ti'].xcom_pull(key='best_model', task_ids='select_best_model')
    model_path = f"/models/{best_model}_{context['ds']}.pkl"

    # Copy to production location
    import shutil
    shutil.copy(model_path, "/production/model.pkl")

    # Get run_id from best model's training task
    best_model_type = context['ti'].xcom_pull(key='best_model', task_ids='select_best_model')

    # Determine task_id based on model type
    if best_model_type == 'rf':
        train_task_id = 'train_models.train_random_forest'
    else:  # lr
        train_task_id = 'train_models.train_logistic_regression'

    run_id = context['ti'].xcom_pull(
        key=f'{best_model_type}_run_id',
        task_ids=train_task_id
    )

    # Update model registry with actual run_id
    import mlflow
    mlflow.register_model(f"runs:/{run_id}/model", "production-model")

    logger.info(f"Deployed {best_model} to production")

deploy = PythonOperator(
    task_id='deploy_model',
    python_callable=deploy_model,
    dag=dag
)

# Define dependencies
validate >> features >> train_group >> select >> deploy

Dynamic DAG Generation

# dags/dynamic_training.py
from airflow import DAG
from airflow.operators.python import PythonOperator
from datetime import datetime

# Configuration for multiple models
MODELS_CONFIG = {
    'fraud_detection': {
        'data_source': 's3://data/fraud/',
        'features': ['amount', 'merchant', 'time'],
        'schedule': '0 3 * * *'
    },
    'churn_prediction': {
        'data_source': 's3://data/churn/',
        'features': ['usage', 'tenure', 'support_calls'],
        'schedule': '0 4 * * *'
    }
}

def create_training_dag(model_name, config):
    """Generate DAG for a specific model."""

    dag = DAG(
        f'train_{model_name}',
        schedule_interval=config['schedule'],
        start_date=datetime(2025, 1, 1),
        catchup=False,
        tags=['ml', model_name]
    )

    def train(**context):
        logger.info(f"Training {model_name} with {config['features']}")
        # Training logic here

    train_task = PythonOperator(
        task_id='train',
        python_callable=train,
        dag=dag
    )

    return dag

# Generate DAG for each model
for model_name, config in MODELS_CONFIG.items():
    globals()[f'train_{model_name}'] = create_training_dag(model_name, config)

Error Handling and Retries

from airflow.exceptions import AirflowException
from airflow.utils.email import send_email

def task_with_retry(**context):
    """Task with custom retry logic."""
    try:
        # Task logic
        result = risky_operation()

    except TemporaryError as e:
        # Retry for temporary errors
        logger.warning(f"Temporary error, will retry: {e}")
        raise AirflowException("Retrying due to temporary error")

    except PermanentError as e:
        # Don't retry permanent errors
        logger.error(f"Permanent error: {e}")
        send_alert_email(str(e))
        raise AirflowException("Permanent failure") from e

def on_failure_callback(context):
    """Custom failure handling."""
    ti = context['task_instance']

    send_email(
        to=['ml-team@example.com'],
        subject=f"Task Failed: {ti.task_id}",
        html_content=f"""
        <h3>Task Failure</h3>
        <p>Task: {ti.task_id}</p>
        <p>DAG: {ti.dag_id}</p>
        <p>Execution Date: {context['ds']}</p>
        <p>Error: {context.get('exception')}</p>
        """
    )

task = PythonOperator(
    task_id='risky_task',
    python_callable=task_with_retry,
    on_failure_callback=on_failure_callback,
    dag=dag
)

Sensors for Data Availability

from airflow.sensors.filesystem import FileSensor
from airflow.sensors.python import PythonSensor
from airflow.sensors.external_task import ExternalTaskSensor

# Wait for file to appear
wait_for_data = FileSensor(
    task_id='wait_for_data_file',
    filepath='/data/input/{{ ds }}.csv',
    poke_interval=60,  # Check every 60 seconds
    timeout=3600,  # Timeout after 1 hour
    mode='reschedule',  # Don't block worker slot
    dag=dag
)

# Wait for upstream DAG
wait_for_upstream = ExternalTaskSensor(
    task_id='wait_for_data_pipeline',
    external_dag_id='data_ingestion',
    external_task_id='export_data',
    execution_delta=timedelta(hours=1),
    dag=dag
)

# Custom condition sensor
def check_data_quality():
    """Check if data meets quality threshold."""
    import pandas as pd
    df = pd.read_csv('/data/latest.csv')
    return len(df) > 1000 and df.isnull().sum().sum() == 0

quality_sensor = PythonSensor(
    task_id='check_data_quality',
    python_callable=check_data_quality,
    poke_interval=300,
    timeout=7200,
    dag=dag
)

Best Practices

1. Idempotency: Tasks should produce same result when re-run 2. XCom for small data only: Use external storage (S3, DB) for large data 3. Task timeouts: Set reasonable execution_timeout 4. Connection pooling: Reuse database connections 5. Monitoring: Use Airflow UI, logs, and custom metrics

Kubeflow Pipelines and MLflow Integration

Production patterns for Kubeflow Pipelines (KFP) orchestration and MLflow experiment tracking in ML pipelines.

Kubeflow Pipelines Components

Component Creation Pattern

from kfp import dsl
from kfp.dsl import component, Input, Output, Dataset, Model, Metrics
from typing import NamedTuple

@component(
    base_image='python:3.11',
    packages_to_install=['pandas==2.1.0', 'scikit-learn==1.3.0']
)
def load_data(
    data_path: str,
    output_dataset: Output[Dataset]
):
    """Load and validate dataset."""
    import pandas as pd
    import logging

    logging.basicConfig(level=logging.INFO)
    logger = logging.getLogger(__name__)

    try:
        # Load data
        df = pd.read_csv(data_path)

        # Validation
        assert len(df) > 100, f"Insufficient data: {len(df)} rows"
        assert df.isnull().sum().sum() < len(df) * 0.1, "Too many nulls"

        # Save output
        df.to_csv(output_dataset.path, index=False)

        logger.info(f"Loaded {len(df)} rows, {len(df.columns)} columns")

    except Exception as e:
        logger.error(f"Data loading failed: {e}")
        raise


@component(
    base_image='python:3.11',
    packages_to_install=['pandas==2.1.0', 'scikit-learn==1.3.0', 'numpy==1.24.0']
)
def preprocess_data(
    input_dataset: Input[Dataset],
    output_dataset: Output[Dataset],
    train_split: float = 0.8
) -> NamedTuple('Outputs', [('num_train', int), ('num_test', int)]):
    """Preprocess and split data."""
    import pandas as pd
    import numpy as np
    from sklearn.model_selection import train_test_split
    from sklearn.preprocessing import StandardScaler
    import logging
    from collections import namedtuple

    logger = logging.getLogger(__name__)

    # Load data
    df = pd.read_csv(input_dataset.path)

    # Feature engineering
    df['feature_ratio'] = df['feature_a'] / (df['feature_b'] + 1e-6)
    df['feature_log'] = np.log1p(df['feature_c'])

    # Split
    X = df.drop('target', axis=1)
    y = df['target']

    X_train, X_test, y_train, y_test = train_test_split(
        X, y, train_size=train_split, random_state=42
    )

    # Scale features
    scaler = StandardScaler()
    X_train = scaler.fit_transform(X_train)
    X_test = scaler.transform(X_test)

    # Save processed data
    train_df = pd.DataFrame(X_train)
    train_df['target'] = y_train.values

    test_df = pd.DataFrame(X_test)
    test_df['target'] = y_test.values

    combined = pd.concat([train_df, test_df])
    combined.to_csv(output_dataset.path, index=False)

    logger.info(f"Train: {len(train_df)}, Test: {len(test_df)}")

    outputs = namedtuple('Outputs', ['num_train', 'num_test'])
    return outputs(len(train_df), len(test_df))


@component(
    base_image='python:3.11',
    packages_to_install=[
        'pandas==2.1.0',
        'scikit-learn==1.3.0',
        'mlflow==2.8.0',
        'joblib==1.3.0'
    ]
)
def train_model(
    input_dataset: Input[Dataset],
    output_model: Output[Model],
    output_metrics: Output[Metrics],
    model_type: str = 'random_forest',
    n_estimators: int = 100,
    mlflow_tracking_uri: str = 'http://mlflow:5000'
) -> float:
    """Train model with MLflow tracking."""
    import pandas as pd
    from sklearn.ensemble import RandomForestClassifier
    from sklearn.linear_model import LogisticRegression
    from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score
    import mlflow
    import mlflow.sklearn
    import joblib
    import logging

    logger = logging.getLogger(__name__)

    # Setup MLflow
    mlflow.set_tracking_uri(mlflow_tracking_uri)
    mlflow.set_experiment('kfp-ml-pipeline')

    # Load data
    df = pd.read_csv(input_dataset.path)
    X = df.drop('target', axis=1)
    y = df['target']

    # Select model
    if model_type == 'random_forest':
        model = RandomForestClassifier(n_estimators=n_estimators, random_state=42)
    elif model_type == 'logistic_regression':
        model = LogisticRegression(max_iter=1000)
    else:
        raise ValueError(f"Unknown model type: {model_type}")

    # Train with MLflow
    with mlflow.start_run(run_name=f"{model_type}_kfp"):
        # Train
        model.fit(X, y)

        # Evaluate
        y_pred = model.predict(X)
        accuracy = accuracy_score(y, y_pred)
        precision = precision_score(y, y_pred, average='weighted')
        recall = recall_score(y, y_pred, average='weighted')
        f1 = f1_score(y, y_pred, average='weighted')

        # Log to MLflow
        mlflow.log_param('model_type', model_type)
        mlflow.log_param('n_estimators', n_estimators)
        mlflow.log_metric('accuracy', accuracy)
        mlflow.log_metric('precision', precision)
        mlflow.log_metric('recall', recall)
        mlflow.log_metric('f1', f1)

        # Log model
        mlflow.sklearn.log_model(model, 'model')

        # Save model artifact
        joblib.dump(model, output_model.path)

        # Save metrics
        output_metrics.log_metric('accuracy', accuracy)
        output_metrics.log_metric('precision', precision)
        output_metrics.log_metric('recall', recall)
        output_metrics.log_metric('f1', f1)

        logger.info(f"Model trained: {model_type}, Accuracy: {accuracy:.4f}")

        return accuracy

Complete KFP Pipeline

from kfp import dsl

@dsl.pipeline(
    name='ML Training Pipeline',
    description='End-to-end ML pipeline with Kubeflow and MLflow'
)
def ml_training_pipeline(
    data_path: str = 's3://data/train.csv',
    model_type: str = 'random_forest',
    n_estimators: int = 100,
    train_split: float = 0.8,
    mlflow_uri: str = 'http://mlflow:5000'
):
    """Complete ML training pipeline."""

    # Step 1: Load data
    load_task = load_data(data_path=data_path)

    # Step 2: Preprocess
    preprocess_task = preprocess_data(
        input_dataset=load_task.outputs['output_dataset'],
        train_split=train_split
    )

    # Step 3: Train model
    train_task = train_model(
        input_dataset=preprocess_task.outputs['output_dataset'],
        model_type=model_type,
        n_estimators=n_estimators,
        mlflow_tracking_uri=mlflow_uri
    )

    # Step 4: Deploy if accuracy > threshold
    with dsl.Condition(train_task.output > 0.85, name='accuracy-check'):
        deploy_task = deploy_model(
            input_model=train_task.outputs['output_model']
        )


@component(base_image='python:3.11')
def deploy_model(
    input_model: Input[Model],
    deployment_endpoint: str = 'http://model-server:8000'
):
    """Deploy model to production."""
    import shutil
    import requests
    import logging

    logger = logging.getLogger(__name__)

    # Copy model to deployment location
    shutil.copy(input_model.path, '/production/model.pkl')

    # Notify deployment service
    response = requests.post(
        f"{deployment_endpoint}/reload",
        json={'model_path': '/production/model.pkl'}
    )

    if response.status_code == 200:
        logger.info("Model deployed successfully")
    else:
        raise Exception(f"Deployment failed: {response.text}")

Running KFP Pipelines

from kfp import compiler
from kfp.client import Client

# Compile pipeline
compiler.Compiler().compile(
    pipeline_func=ml_training_pipeline,
    package_path='ml_pipeline.yaml'
)

# Submit to Kubeflow
client = Client(host='http://kubeflow.example.com')

run = client.create_run_from_pipeline_func(
    ml_training_pipeline,
    arguments={
        'data_path': 's3://my-bucket/data/train.csv',
        'model_type': 'random_forest',
        'n_estimators': 200,
        'mlflow_uri': 'http://mlflow.example.com'
    },
    experiment_name='ml-training'
)

print(f"Pipeline run created: {run.run_id}")

MLflow Tracking Integration

Experiment Tracking Setup

import mlflow
from mlflow.tracking import MlflowClient

# Configure MLflow
mlflow.set_tracking_uri('http://mlflow.example.com:5000')
mlflow.set_experiment('production-training')

def train_with_mlflow(X_train, y_train, X_test, y_test, config):
    """Train model with comprehensive MLflow tracking."""

    with mlflow.start_run(run_name=config['run_name']):
        # Log parameters
        mlflow.log_params({
            'model_type': config['model_type'],
            'n_estimators': config.get('n_estimators', 100),
            'learning_rate': config.get('learning_rate', 0.1),
            'max_depth': config.get('max_depth', 5),
            'random_state': 42
        })

        # Log dataset info
        mlflow.log_param('train_size', len(X_train))
        mlflow.log_param('test_size', len(X_test))
        mlflow.log_param('n_features', X_train.shape[1])

        # Train model
        model = create_model(config)
        model.fit(X_train, y_train)

        # Evaluate
        train_score = model.score(X_train, y_train)
        test_score = model.score(X_test, y_test)

        # Log metrics
        mlflow.log_metric('train_accuracy', train_score)
        mlflow.log_metric('test_accuracy', test_score)
        mlflow.log_metric('overfit_ratio', train_score / test_score)

        # Log feature importance
        if hasattr(model, 'feature_importances_'):
            feature_importance = dict(zip(
                [f'feature_{i}' for i in range(len(model.feature_importances_))],
                model.feature_importances_
            ))
            mlflow.log_params(feature_importance)

        # Log model
        mlflow.sklearn.log_model(
            model,
            'model',
            registered_model_name=config.get('model_name', 'production-model')
        )

        # Log artifacts
        import matplotlib.pyplot as plt
        from sklearn.metrics import confusion_matrix, ConfusionMatrixDisplay

        # Confusion matrix
        y_pred = model.predict(X_test)
        cm = confusion_matrix(y_test, y_pred)
        disp = ConfusionMatrixDisplay(cm)
        disp.plot()
        plt.savefig('confusion_matrix.png')
        mlflow.log_artifact('confusion_matrix.png')

        return model, mlflow.active_run().info.run_id

Model Registry Integration

from mlflow.tracking import MlflowClient

client = MlflowClient(tracking_uri='http://mlflow.example.com:5000')

def register_model(run_id: str, model_name: str, stage: str = 'Staging'):
    """Register model in MLflow Model Registry."""

    # Get model URI
    model_uri = f"runs:/{run_id}/model"

    # Register model
    model_version = mlflow.register_model(model_uri, model_name)

    # Transition to stage
    client.transition_model_version_stage(
        name=model_name,
        version=model_version.version,
        stage=stage,
        archive_existing_versions=True
    )

    print(f"Model {model_name} version {model_version.version} -> {stage}")

    return model_version


def promote_to_production(model_name: str, version: str):
    """Promote model version to production."""

    # Get current production model
    prod_versions = client.get_latest_versions(model_name, stages=['Production'])

    # Transition new version to production
    client.transition_model_version_stage(
        name=model_name,
        version=version,
        stage='Production',
        archive_existing_versions=True  # Archive old production
    )

    print(f"Promoted {model_name} v{version} to Production")

    # Archive old versions
    for old_version in prod_versions:
        if old_version.version != version:
            client.transition_model_version_stage(
                name=model_name,
                version=old_version.version,
                stage='Archived'
            )


def load_production_model(model_name: str):
    """Load latest production model."""
    import mlflow.pyfunc

    model_uri = f"models:/{model_name}/Production"
    model = mlflow.pyfunc.load_model(model_uri)

    return model

Hyperparameter Tuning with MLflow

from sklearn.model_selection import GridSearchCV
import mlflow

def hyperparameter_tuning_with_mlflow(X_train, y_train, param_grid):
    """Hyperparameter tuning with MLflow tracking."""

    mlflow.set_experiment('hyperparameter-tuning')

    parent_run = mlflow.start_run(run_name='grid_search')

    # Grid search
    from sklearn.ensemble import RandomForestClassifier
    model = RandomForestClassifier(random_state=42)

    grid_search = GridSearchCV(
        model,
        param_grid,
        cv=5,
        scoring='accuracy',
        n_jobs=-1
    )

    grid_search.fit(X_train, y_train)

    # Log best parameters to parent run
    mlflow.log_params(grid_search.best_params_)
    mlflow.log_metric('best_cv_score', grid_search.best_score_)

    # Log each CV result as child run
    for i, params in enumerate(grid_search.cv_results_['params']):
        with mlflow.start_run(run_name=f'fold_{i}', nested=True):
            mlflow.log_params(params)
            mlflow.log_metric('mean_test_score', grid_search.cv_results_['mean_test_score'][i])
            mlflow.log_metric('std_test_score', grid_search.cv_results_['std_test_score'][i])

    mlflow.end_run()

    return grid_search.best_estimator_, grid_search.best_params_

Artifact Management

Versioned Artifact Storage

import mlflow
from pathlib import Path

def log_artifacts_versioned(run_id: str, artifacts_dir: Path):
    """Log versioned artifacts to MLflow."""

    with mlflow.start_run(run_id=run_id):
        # Log data artifacts
        mlflow.log_artifact(artifacts_dir / 'train.csv', 'data')
        mlflow.log_artifact(artifacts_dir / 'test.csv', 'data')

        # Log preprocessing artifacts
        mlflow.log_artifact(artifacts_dir / 'scaler.pkl', 'preprocessing')
        mlflow.log_artifact(artifacts_dir / 'encoder.pkl', 'preprocessing')

        # Log evaluation artifacts
        mlflow.log_artifact(artifacts_dir / 'metrics.json', 'evaluation')
        mlflow.log_artifact(artifacts_dir / 'confusion_matrix.png', 'evaluation')

        # Log model artifacts
        mlflow.log_artifact(artifacts_dir / 'model.pkl', 'model')
        mlflow.log_artifact(artifacts_dir / 'model_metadata.json', 'model')


def download_artifacts(run_id: str, artifact_path: str, dst_path: Path):
    """Download artifacts from MLflow."""

    client = MlflowClient()

    # Download artifacts
    client.download_artifacts(
        run_id=run_id,
        path=artifact_path,
        dst_path=str(dst_path)
    )

    print(f"Downloaded artifacts to {dst_path}")

Best Practices

Component Design

1. Single Responsibility: Each component does one thing 2. Type Hints: Use Input/Output types for data passing 3. Error Handling: Comprehensive try/except in components 4. Logging: Log all important events 5. Versioning: Pin package versions in base_image

MLflow Tracking

1. Experiment Organization: One experiment per project/model type 2. Run Naming: Descriptive run names with timestamp 3. Parameter Logging: Log all hyperparameters 4. Metric Logging: Log train AND test metrics 5. Artifact Logging: Save models, plots, data samples

Pipeline Orchestration

1. Conditional Execution: Use dsl.Condition for branching 2. Parallel Execution: Use dsl.ParallelFor for batch jobs 3. Resource Limits: Set CPU/memory limits on components 4. Retry Logic: Configure retries for flaky components 5. Monitoring: Use KFP UI to monitor pipeline health

Pipeline Monitoring and Data Quality

Production monitoring for ML pipelines including data quality checks, drift detection, and alert configuration.

Data Quality Validation

Schema Validation

from typing import Dict, List, Optional
import pandas as pd
from dataclasses import dataclass
import logging

logger = logging.getLogger(__name__)


@dataclass
class ColumnSchema:
    """Schema definition for a column."""
    name: str
    dtype: str
    nullable: bool = True
    min_value: Optional[float] = None
    max_value: Optional[float] = None
    allowed_values: Optional[List] = None


class DataValidator:
    """Validate data quality against schema."""

    def __init__(self, schema: List[ColumnSchema]):
        self.schema = {col.name: col for col in schema}

    def validate(self, df: pd.DataFrame) -> tuple[bool, List[str]]:
        """
        Validate DataFrame against schema.

        Returns:
            (is_valid: bool, errors: List[str])
        """
        errors = []

        # Check columns exist
        expected_cols = set(self.schema.keys())
        actual_cols = set(df.columns)

        missing = expected_cols - actual_cols
        extra = actual_cols - expected_cols

        if missing:
            errors.append(f"Missing columns: {missing}")
        if extra:
            errors.append(f"Extra columns: {extra}")

        # Validate each column
        for col_name, col_schema in self.schema.items():
            if col_name not in df.columns:
                continue

            col = df[col_name]

            # Check dtype
            if str(col.dtype) != col_schema.dtype:
                errors.append(
                    f"Column {col_name}: expected dtype {col_schema.dtype}, "
                    f"got {col.dtype}"
                )

            # Check nulls
            null_count = col.isnull().sum()
            if null_count > 0 and not col_schema.nullable:
                errors.append(
                    f"Column {col_name}: contains {null_count} nulls "
                    f"(not nullable)"
                )

            # Check numeric ranges
            if col_schema.min_value is not None:
                if col.min() < col_schema.min_value:
                    errors.append(
                        f"Column {col_name}: min value {col.min()} "
                        f"below {col_schema.min_value}"
                    )

            if col_schema.max_value is not None:
                if col.max() > col_schema.max_value:
                    errors.append(
                        f"Column {col_name}: max value {col.max()} "
                        f"above {col_schema.max_value}"
                    )

            # Check allowed values
            if col_schema.allowed_values is not None:
                invalid = set(col.unique()) - set(col_schema.allowed_values)
                if invalid:
                    errors.append(
                        f"Column {col_name}: invalid values {invalid}"
                    )

        is_valid = len(errors) == 0
        return is_valid, errors


# Example usage
schema = [
    ColumnSchema(name='user_id', dtype='int64', nullable=False),
    ColumnSchema(name='age', dtype='int64', min_value=0, max_value=120),
    ColumnSchema(
        name='status',
        dtype='object',
        allowed_values=['active', 'inactive', 'pending']
    ),
    ColumnSchema(name='score', dtype='float64', min_value=0.0, max_value=1.0),
]

validator = DataValidator(schema)
is_valid, errors = validator.validate(df)

if not is_valid:
    logger.error(f"Data validation failed: {errors}")
    raise ValueError("Data quality check failed")

Statistical Quality Checks

import numpy as np
from scipy import stats


class StatisticalValidator:
    """Statistical data quality checks."""

    def __init__(self, reference_df: pd.DataFrame):
        """
        Args:
            reference_df: Historical data for comparison
        """
        self.reference_df = reference_df
        self.reference_stats = self._compute_stats(reference_df)

    def _compute_stats(self, df: pd.DataFrame) -> Dict:
        """Compute reference statistics."""
        stats_dict = {}

        for col in df.select_dtypes(include=[np.number]).columns:
            stats_dict[col] = {
                'mean': df[col].mean(),
                'std': df[col].std(),
                'min': df[col].min(),
                'max': df[col].max(),
                'median': df[col].median(),
                'q25': df[col].quantile(0.25),
                'q75': df[col].quantile(0.75)
            }

        return stats_dict

    def validate_distribution(
        self,
        df: pd.DataFrame,
        threshold: float = 0.05
    ) -> tuple[bool, Dict]:
        """
        Check if distributions match reference using KS test.

        Args:
            df: New data to validate
            threshold: p-value threshold (default 0.05)

        Returns:
            (is_valid: bool, results: Dict)
        """
        results = {}

        for col in df.select_dtypes(include=[np.number]).columns:
            if col not in self.reference_stats:
                continue

            # Kolmogorov-Smirnov test
            statistic, p_value = stats.ks_2samp(
                self.reference_df[col].dropna(),
                df[col].dropna()
            )

            is_valid = p_value > threshold

            results[col] = {
                'ks_statistic': statistic,
                'p_value': p_value,
                'is_valid': is_valid,
                'reference_mean': self.reference_stats[col]['mean'],
                'current_mean': df[col].mean(),
                'mean_shift': abs(df[col].mean() - self.reference_stats[col]['mean'])
            }

        overall_valid = all(r['is_valid'] for r in results.values())
        return overall_valid, results

    def check_outliers(
        self,
        df: pd.DataFrame,
        max_outlier_pct: float = 0.05
    ) -> tuple[bool, Dict]:
        """
        Check for excessive outliers using IQR method.

        Args:
            df: Data to check
            max_outlier_pct: Max percentage of outliers allowed

        Returns:
            (is_valid: bool, outlier_info: Dict)
        """
        outlier_info = {}

        for col in df.select_dtypes(include=[np.number]).columns:
            Q1 = df[col].quantile(0.25)
            Q3 = df[col].quantile(0.75)
            IQR = Q3 - Q1

            lower_bound = Q1 - 1.5 * IQR
            upper_bound = Q3 + 1.5 * IQR

            outliers = ((df[col] < lower_bound) | (df[col] > upper_bound)).sum()
            outlier_pct = outliers / len(df)

            outlier_info[col] = {
                'outlier_count': outliers,
                'outlier_pct': outlier_pct,
                'is_valid': outlier_pct <= max_outlier_pct,
                'bounds': (lower_bound, upper_bound)
            }

        overall_valid = all(info['is_valid'] for info in outlier_info.values())
        return overall_valid, outlier_info

Data Drift Detection

Distribution Drift Monitor

from typing import Callable
import json
from datetime import datetime


class DriftMonitor:
    """Monitor data drift over time."""

    def __init__(
        self,
        reference_data: pd.DataFrame,
        alert_callback: Optional[Callable] = None,
        drift_threshold: float = 0.1
    ):
        self.reference_data = reference_data
        self.alert_callback = alert_callback
        self.drift_threshold = drift_threshold
        self.drift_history = []

    def detect_drift(self, current_data: pd.DataFrame) -> Dict:
        """
        Detect drift in current data.

        Returns:
            Dictionary with drift scores and alerts
        """
        drift_scores = {}
        alerts = []

        # Numerical features
        for col in current_data.select_dtypes(include=[np.number]).columns:
            if col not in self.reference_data.columns:
                continue

            # KS test
            statistic, p_value = stats.ks_2samp(
                self.reference_data[col].dropna(),
                current_data[col].dropna()
            )

            drift_scores[col] = {
                'ks_statistic': statistic,
                'p_value': p_value,
                'drifted': statistic > self.drift_threshold
            }

            if statistic > self.drift_threshold:
                alerts.append({
                    'column': col,
                    'type': 'distribution_drift',
                    'severity': 'high' if statistic > 0.2 else 'medium',
                    'score': statistic,
                    'timestamp': datetime.now().isoformat()
                })

        # Categorical features
        for col in current_data.select_dtypes(include=['object', 'category']).columns:
            if col not in self.reference_data.columns:
                continue

            ref_dist = self.reference_data[col].value_counts(normalize=True)
            curr_dist = current_data[col].value_counts(normalize=True)

            # Chi-squared test
            try:
                chi2, p_value = stats.chisquare(
                    curr_dist.reindex(ref_dist.index, fill_value=0),
                    ref_dist
                )

                drift_scores[col] = {
                    'chi2_statistic': chi2,
                    'p_value': p_value,
                    'drifted': p_value < 0.05
                }

                if p_value < 0.05:
                    alerts.append({
                        'column': col,
                        'type': 'categorical_drift',
                        'severity': 'high' if p_value < 0.01 else 'medium',
                        'chi2': chi2,
                        'timestamp': datetime.now().isoformat()
                    })
            except ValueError as e:
                logger.warning(f"Chi-squared test failed for {col}: {e}")

        # Store drift history
        self.drift_history.append({
            'timestamp': datetime.now().isoformat(),
            'drift_scores': drift_scores,
            'num_alerts': len(alerts)
        })

        # Trigger alerts
        if alerts and self.alert_callback:
            self.alert_callback(alerts)

        return {
            'drift_detected': len(alerts) > 0,
            'drift_scores': drift_scores,
            'alerts': alerts
        }

Alert Configuration

Alert Manager

import requests
from typing import List, Dict
import smtplib
from email.mime.text import MIMEText
from email.mime.multipart import MIMEMultipart


class AlertManager:
    """Send alerts for pipeline failures and drift."""

    def __init__(
        self,
        slack_webhook: Optional[str] = None,
        email_config: Optional[Dict] = None
    ):
        self.slack_webhook = slack_webhook
        self.email_config = email_config

    def send_alert(
        self,
        title: str,
        message: str,
        severity: str = 'info',
        details: Optional[Dict] = None
    ):
        """Send alert through configured channels."""

        # Format alert
        alert = {
            'title': title,
            'message': message,
            'severity': severity,
            'timestamp': datetime.now().isoformat(),
            'details': details or {}
        }

        # Send to Slack
        if self.slack_webhook:
            self._send_slack_alert(alert)

        # Send email
        if self.email_config:
            self._send_email_alert(alert)

        # Log
        logger.warning(f"Alert: {title} - {message}")

    def _send_slack_alert(self, alert: Dict):
        """Send alert to Slack."""
        emoji = {
            'critical': ':rotating_light:',
            'high': ':warning:',
            'medium': ':exclamation:',
            'low': ':information_source:',
            'info': ':mega:'
        }.get(alert['severity'], ':mega:')

        message = {
            'text': f"{emoji} *{alert['title']}*",
            'blocks': [
                {
                    'type': 'section',
                    'text': {
                        'type': 'mrkdwn',
                        'text': f"{emoji} *{alert['title']}*\n{alert['message']}"
                    }
                },
                {
                    'type': 'context',
                    'elements': [
                        {
                            'type': 'mrkdwn',
                            'text': f"Severity: `{alert['severity']}` | Time: {alert['timestamp']}"
                        }
                    ]
                }
            ]
        }

        if alert['details']:
            message['blocks'].append({
                'type': 'section',
                'text': {
                    'type': 'mrkdwn',
                    'text': f"```{json.dumps(alert['details'], indent=2)}```"
                }
            })

        try:
            response = requests.post(
                self.slack_webhook,
                json=message,
                timeout=10
            )
            response.raise_for_status()
        except Exception as e:
            logger.error(f"Failed to send Slack alert: {e}")

    def _send_email_alert(self, alert: Dict):
        """Send alert via email."""
        msg = MIMEMultipart()
        msg['From'] = self.email_config['from']
        msg['To'] = ', '.join(self.email_config['to'])
        msg['Subject'] = f"[{alert['severity'].upper()}] {alert['title']}"

        body = f"""
        Alert: {alert['title']}
        Severity: {alert['severity']}
        Time: {alert['timestamp']}

        {alert['message']}

        Details:
        {json.dumps(alert['details'], indent=2)}
        """

        msg.attach(MIMEText(body, 'plain'))

        try:
            with smtplib.SMTP(
                self.email_config['smtp_host'],
                self.email_config['smtp_port']
            ) as server:
                server.starttls()
                server.login(
                    self.email_config['username'],
                    self.email_config['password']
                )
                server.send_message(msg)
        except Exception as e:
            logger.error(f"Failed to send email alert: {e}")

    def send_drift_alert(self, drift_results: Dict):
        """Send drift detection alert."""
        alerts = drift_results.get('alerts', [])

        if not alerts:
            return

        # Group by severity
        critical = [a for a in alerts if a['severity'] == 'critical']
        high = [a for a in alerts if a['severity'] == 'high']
        medium = [a for a in alerts if a['severity'] == 'medium']

        message = f"Data drift detected in {len(alerts)} features\n"
        if critical:
            message += f"- Critical: {len(critical)}\n"
        if high:
            message += f"- High: {len(high)}\n"
        if medium:
            message += f"- Medium: {len(medium)}\n"

        self.send_alert(
            title='Data Drift Detected',
            message=message,
            severity='high' if (critical or high) else 'medium',
            details=drift_results
        )

Pipeline Health Monitoring

Pipeline Metrics Tracker

from prometheus_client import Counter, Histogram, Gauge, generate_latest
from flask import Flask, Response

app = Flask(__name__)

# Metrics
pipeline_runs = Counter(
    'pipeline_runs_total',
    'Total pipeline runs',
    ['pipeline_name', 'status']
)

pipeline_duration = Histogram(
    'pipeline_duration_seconds',
    'Pipeline execution time',
    ['pipeline_name'],
    buckets=[60, 300, 600, 1800, 3600, 7200]  # 1min to 2 hours
)

data_quality_score = Gauge(
    'data_quality_score',
    'Data quality score (0-1)',
    ['pipeline_name']
)

drift_score = Gauge(
    'drift_score',
    'Data drift score',
    ['pipeline_name', 'feature']
)


class PipelineMetrics:
    """Track pipeline metrics."""

    @staticmethod
    def record_run(pipeline_name: str, status: str, duration: float):
        """Record pipeline run."""
        pipeline_runs.labels(
            pipeline_name=pipeline_name,
            status=status
        ).inc()

        pipeline_duration.labels(
            pipeline_name=pipeline_name
        ).observe(duration)

    @staticmethod
    def record_quality(pipeline_name: str, score: float):
        """Record data quality score."""
        data_quality_score.labels(
            pipeline_name=pipeline_name
        ).set(score)

    @staticmethod
    def record_drift(pipeline_name: str, feature: str, score: float):
        """Record drift score."""
        drift_score.labels(
            pipeline_name=pipeline_name,
            feature=feature
        ).set(score)


@app.route('/metrics')
def metrics():
    """Expose Prometheus metrics."""
    return Response(generate_latest(), mimetype='text/plain')


if __name__ == '__main__':
    app.run(host='0.0.0.0', port=9090)

Airflow Integration

Data Quality Airflow Operator

from airflow.models import BaseOperator
from airflow.utils.decorators import apply_defaults


class DataQualityOperator(BaseOperator):
    """Airflow operator for data quality checks."""

    @apply_defaults
    def __init__(
        self,
        data_path: str,
        schema: List[ColumnSchema],
        drift_threshold: float = 0.1,
        reference_data_path: Optional[str] = None,
        *args,
        **kwargs
    ):
        super().__init__(*args, **kwargs)
        self.data_path = data_path
        self.schema = schema
        self.drift_threshold = drift_threshold
        self.reference_data_path = reference_data_path

    def execute(self, context):
        """Execute data quality checks."""
        import pandas as pd

        # Load data
        df = pd.read_csv(self.data_path)

        # Schema validation
        validator = DataValidator(self.schema)
        is_valid, errors = validator.validate(df)

        if not is_valid:
            raise ValueError(f"Schema validation failed: {errors}")

        # Statistical validation
        if self.reference_data_path:
            reference_df = pd.read_csv(self.reference_data_path)
            stat_validator = StatisticalValidator(reference_df)

            dist_valid, dist_results = stat_validator.validate_distribution(df)
            if not dist_valid:
                self.log.warning(f"Distribution drift detected: {dist_results}")

            outlier_valid, outlier_info = stat_validator.check_outliers(df)
            if not outlier_valid:
                self.log.warning(f"Excessive outliers: {outlier_info}")

        self.log.info("Data quality checks passed")
        return True

Best Practices

1. Comprehensive Validation: Check schema, statistics, and distributions 2. Automated Alerts: Configure alerts for critical issues 3. Track History: Store validation results for trend analysis 4. Gradual Rollout: Don't fail pipelines on first drift detection 5. Actionable Metrics: Track metrics that drive decisions 6. Alert Fatigue: Set appropriate thresholds to avoid noise 7. Documentation: Document all quality checks and thresholds