ML Pipeline Automation

ML Pipeline Automation

ML pipeline automation orchestrates the entire machine learning workflow from data ingestion through model deployment, ensuring reproducibility, scalability, and reliability.

Pipeline Components

• Data Ingestion: Collecting data from multiple sources
• Data Processing: Cleaning, transformation, feature engineering
• Model Training: Training and hyperparameter tuning
• Validation: Cross-validation and testing
• Deployment: Moving models to production
• Monitoring: Tracking performance metrics

Orchestration Platforms

• Apache Airflow: Workflow scheduling with DAGs
• Kubeflow: Kubernetes-native ML workflows
• Jenkins: CI/CD for ML pipelines
• Prefect: Modern data flow orchestration
• Dagster: Asset-driven orchestration

Python Implementation

import pandas as pd
import numpy as np
from sklearn.datasets import make_classification
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score, f1_score
import joblib
import logging
from datetime import datetime
import json
import os

# Airflow imports
from airflow import DAG
from airflow.operators.python import PythonOperator
from airflow.operators.bash import BashOperator
from airflow.utils.dates import days_ago

# MLflow for tracking
import mlflow
import mlflow.sklearn

# Logging setup
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

print("=== 1. Modular Pipeline Functions ===")

# Data ingestion
def ingest_data(**context):
    """Ingest and load data"""
    logger.info("Starting data ingestion...")

    X, y = make_classification(n_samples=2000, n_features=30,
                              n_informative=20, random_state=42)
    data = pd.DataFrame(X, columns=[f'feature_{i}' for i in range(X.shape[1])])
    data['target'] = y

    # Save to disk
    data_path = '/tmp/raw_data.csv'
    data.to_csv(data_path, index=False)

    context['task_instance'].xcom_push(key='data_path', value=data_path)
    logger.info(f"Data ingested: {len(data)} rows")
    return {'status': 'success', 'samples': len(data)}

# Data processing
def process_data(**context):
    """Clean and preprocess data"""
    logger.info("Starting data processing...")

    # Get data path from previous task
    task_instance = context['task_instance']
    data_path = task_instance.xcom_pull(key='data_path', task_ids='ingest_data')

    data = pd.read_csv(data_path)

    # Handle missing values
    data = data.fillna(data.mean())

    # Remove duplicates
    data = data.drop_duplicates()

    # Remove outliers (simple approach)
    numeric_cols = data.select_dtypes(include=[np.number]).columns
    for col in numeric_cols:
        Q1 = data[col].quantile(0.25)
        Q3 = data[col].quantile(0.75)
        IQR = Q3 - Q1
        data = data[(data[col] >= Q1 - 1.5 * IQR) & (data[col] <= Q3 + 1.5 * IQR)]

    processed_path = '/tmp/processed_data.csv'
    data.to_csv(processed_path, index=False)

    task_instance.xcom_push(key='processed_path', value=processed_path)
    logger.info(f"Data processed: {len(data)} rows after cleaning")
    return {'status': 'success', 'rows_remaining': len(data)}

# Feature engineering
def engineer_features(**context):
    """Create new features"""
    logger.info("Starting feature engineering...")

    task_instance = context['task_instance']
    processed_path = task_instance.xcom_pull(key='processed_path', task_ids='process_data')

    data = pd.read_csv(processed_path)

    # Create interaction features
    feature_cols = [col for col in data.columns if col.startswith('feature_')]
    for i in range(min(5, len(feature_cols))):
        for j in range(i+1, min(6, len(feature_cols))):
            data[f'interaction_{i}_{j}'] = data[feature_cols[i]] * data[feature_cols[j]]

    # Create polynomial features
    for col in feature_cols[:5]:
        data[f'{col}_squared'] = data[col] ** 2

    engineered_path = '/tmp/engineered_data.csv'
    data.to_csv(engineered_path, index=False)

    task_instance.xcom_push(key='engineered_path', value=engineered_path)
    logger.info(f"Features engineered: {len(data.columns)} total features")
    return {'status': 'success', 'features': len(data.columns)}

# Train model
def train_model(**context):
    """Train ML model"""
    logger.info("Starting model training...")

    task_instance = context['task_instance']
    engineered_path = task_instance.xcom_pull(key='engineered_path', task_ids='engineer_features')

    data = pd.read_csv(engineered_path)

    X = data.drop('target', axis=1)
    y = data['target']

    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

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

    # Train model
    model = RandomForestClassifier(n_estimators=100, max_depth=15, random_state=42)
    model.fit(X_train_scaled, y_train)

    # Evaluate
    y_pred = model.predict(X_test_scaled)
    accuracy = accuracy_score(y_test, y_pred)
    f1 = f1_score(y_test, y_pred)

    # Save model
    model_path = '/tmp/model.pkl'
    scaler_path = '/tmp/scaler.pkl'
    joblib.dump(model, model_path)
    joblib.dump(scaler, scaler_path)

    task_instance.xcom_push(key='model_path', value=model_path)
    task_instance.xcom_push(key='scaler_path', value=scaler_path)

    # Log to MLflow
    with mlflow.start_run():
        mlflow.log_param('n_estimators', 100)
        mlflow.log_param('max_depth', 15)
        mlflow.log_metric('accuracy', accuracy)
        mlflow.log_metric('f1_score', f1)
        mlflow.sklearn.log_model(model, 'model')

    logger.info(f"Model trained: Accuracy={accuracy:.4f}, F1={f1:.4f}")
    return {'status': 'success', 'accuracy': accuracy, 'f1_score': f1}

# Validate model
def validate_model(**context):
    """Validate model performance"""
    logger.info("Starting model validation...")

    task_instance = context['task_instance']
    model_path = task_instance.xcom_pull(key='model_path', task_ids='train_model')
    engineered_path = task_instance.xcom_pull(key='engineered_path', task_ids='engineer_features')

    model = joblib.load(model_path)
    data = pd.read_csv(engineered_path)

    X = data.drop('target', axis=1)
    y = data['target']

    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

    scaler_path = task_instance.xcom_pull(key='scaler_path', task_ids='train_model')
    scaler = joblib.load(scaler_path)
    X_test_scaled = scaler.transform(X_test)

    # Validate
    y_pred = model.predict(X_test_scaled)
    accuracy = accuracy_score(y_test, y_pred)

    validation_result = {
        'status': 'success' if accuracy > 0.85 else 'failed',
        'accuracy': accuracy,
        'threshold': 0.85,
        'timestamp': datetime.now().isoformat()
    }

    task_instance.xcom_push(key='validation_result', value=json.dumps(validation_result))

    logger.info(f"Validation result: {validation_result}")
    return validation_result

# Deploy model
def deploy_model(**context):
    """Deploy validated model"""
    logger.info("Starting model deployment...")

    task_instance = context['task_instance']
    validation_result = json.loads(task_instance.xcom_pull(
        key='validation_result', task_ids='validate_model'))

    if validation_result['status'] != 'success':
        logger.warning("Validation failed, deployment skipped")
        return {'status': 'skipped', 'reason': 'validation_failed'}

    model_path = task_instance.xcom_pull(key='model_path', task_ids='train_model')
    scaler_path = task_instance.xcom_pull(key='scaler_path', task_ids='train_model')

    # Simulate deployment
    deploy_path = '/tmp/deployed_model/'
    os.makedirs(deploy_path, exist_ok=True)

    import shutil