Serverless MLOps Architecture

import json
import boto3
import pickle
import numpy as np
import pandas as pd
from typing import Dict, Any, List, Optional
import os
import logging
from datetime import datetime
import base64
import io

# AWS SDK clients
s3_client = boto3.client('s3')
sagemaker_client = boto3.client('sagemaker')
cloudwatch = boto3.client('cloudwatch')
stepfunctions = boto3.client('stepfunctions')

logger = logging.getLogger()
logger.setLevel(logging.INFO)

# ============================================================================
# Lambda Function 1: Data Processing Trigger
# ============================================================================
def data_processing_handler(event, context):
    """
    Triggered when new data arrives in S3
    Processes data and triggers training if needed
    """
    try:
        # Parse S3 event
        for record in event['Records']:
            bucket = record['s3']['bucket']['name']
            key = record['s3']['object']['key']
            
            logger.info(f"Processing file: s3://{bucket}/{key}")
            
            # Download and process data
            processed_data = process_incoming_data(bucket, key)
            
            # Check if retraining is needed
            if should_retrain_model(processed_data):
                trigger_training_pipeline(processed_data)
            
            # Update data metrics
            publish_data_metrics(processed_data)
        
        return {
            'statusCode': 200,
            'body': json.dumps({'message': 'Data processed successfully'})
        }
        
    except Exception as e:
        logger.error(f"Data processing failed: {str(e)}")
        return {
            'statusCode': 500,
            'body': json.dumps({'error': str(e)})
        }

def process_incoming_data(bucket: str, key: str) -> Dict[str, Any]:
    """Process raw data from S3"""
    
    # Download file from S3
    response = s3_client.get_object(Bucket=bucket, Key=key)
    data = pd.read_csv(io.BytesIO(response['Body'].read()))
    
    # Data processing pipeline
    processed_data = {
        'raw_rows': len(data),
        'processed_rows': 0,
        'quality_score': 0.0,
        'feature_drift': 0.0,
        's3_location': f's3://{bucket}/{key}'
    }
    
    # Data cleaning
    data_clean = data.dropna()
    processed_data['processed_rows'] = len(data_clean)
    
    # Data quality assessment
    processed_data['quality_score'] = len(data_clean) / len(data)
    
    # Feature drift detection (simplified)
    if 'target' in data_clean.columns:
        current_mean = data_clean['target'].mean()
        historical_mean = get_historical_target_mean()
        processed_data['feature_drift'] = abs(current_mean - historical_mean) / historical_mean
    
    # Save processed data
    processed_key = key.replace('.csv', '_processed.csv')
    csv_buffer = io.StringIO()
    data_clean.to_csv(csv_buffer, index=False)
    
    s3_client.put_object(
        Bucket=bucket,
        Key=processed_key,
        Body=csv_buffer.getvalue()
    )
    
    processed_data['processed_s3_location'] = f's3://{bucket}/{processed_key}'
    
    return processed_data

def should_retrain_model(processed_data: Dict[str, Any]) -> bool:
    """Determine if model retraining is needed"""
    
    # Retrain if data quality is good and there's significant drift
    quality_threshold = 0.8
    drift_threshold = 0.15
    
    return (
        processed_data['quality_score'] > quality_threshold and
        processed_data['feature_drift'] > drift_threshold
    )

def get_historical_target_mean() -> float:
    """Get historical target mean from parameter store or database"""
    # In production, retrieve from AWS Parameter Store or DynamoDB
    return 0.5  # Placeholder

# ============================================================================
# Lambda Function 2: Model Training Orchestrator  
# ============================================================================
def training_orchestrator_handler(event, context):
    """
    Orchestrates model training using SageMaker and Step Functions
    """
    try:
        processed_data = event.get('processed_data', {})
        
        # Create training job configuration
        training_config = create_training_config(processed_data)
        
        # Start Step Functions workflow for training
        workflow_arn = os.environ['TRAINING_WORKFLOW_ARN']
        execution_name = f"training-{datetime.now().strftime('%Y%m%d-%H%M%S')}"
        
        response = stepfunctions.start_execution(
            stateMachineArn=workflow_arn,
            name=execution_name,
            input=json.dumps(training_config)
        )
        
        logger.info(f"Started training workflow: {response['executionArn']}")
        
        return {
            'statusCode': 200,
            'body': json.dumps({
                'execution_arn': response['executionArn'],
                'training_config': training_config
            })
        }
        
    except Exception as e:
        logger.error(f"Training orchestration failed: {str(e)}")
        return {
            'statusCode': 500,
            'body': json.dumps({'error': str(e)})
        }

def create_training_config(processed_data: Dict[str, Any]) -> Dict[str, Any]:
    """Create SageMaker training job configuration"""
    
    timestamp = datetime.now().strftime('%Y-%m-%d-%H-%M-%S')
    
    return {
        'training_job_name': f'ml-model-{timestamp}',
        'algorithm_specification': {
            'training_image': f'{os.environ["ACCOUNT_ID"]}.dkr.ecr.{os.environ["AWS_REGION"]}.amazonaws.com/ml-training:latest',
            'training_input_mode': 'File'
        },
        'input_data_config': [
            {
                'channel_name': 'training',
                'data_source': {
                    'S3DataSource': {
                        'S3DataType': 'S3Prefix',
                        'S3Uri': processed_data['processed_s3_location'],
                        'S3DataDistributionType': 'FullyReplicated'
                    }
                }
            }
        ],
        'output_data_config': {
            'S3OutputPath': f's3://{os.environ["MODEL_BUCKET"]}/models/'
        },
        'resource_config': {
            'InstanceType': 'ml.m5.large',
            'InstanceCount': 1,
            'VolumeSizeInGB': 30
        },
        'stopping_condition': {
            'MaxRuntimeInSeconds': 3600
        },
        'role_arn': os.environ['SAGEMAKER_EXECUTION_ROLE']
    }

def trigger_training_pipeline(processed_data: Dict[str, Any]):
    """Trigger the training pipeline"""
    
    # Invoke training orchestrator Lambda
    lambda_client = boto3.client('lambda')
    
    payload = {
        'processed_data': processed_data
    }
    
    lambda_client.invoke(
        FunctionName=os.environ['TRAINING_ORCHESTRATOR_FUNCTION'],
        InvocationType='Event',  # Async invocation
        Payload=json.dumps(payload)
    )

# ============================================================================
# Lambda Function 3: Model Inference
# ============================================================================
def inference_handler(event, context):
    """
    Serverless model inference endpoint
    Handles both batch and real-time predictions
    """
    try:
        # Handle different event sources
        if 'body' in event:
            # API Gateway request
            request_data = json.loads(event['body'])
        else:
            # Direct Lambda invocation
            request_data = event
        
        # Load model (with caching for warm starts)
        model = load_model_cached()
        
        # Make predictions
        if 'batch' in request_data:
            predictions = batch_inference(model, request_data['batch'])
        else:
            prediction = single_inference(model, request_data['features'])
            predictions = [prediction]
        
        # Log inference metrics
        log_inference_metrics(len(predictions))
        
        return {
            'statusCode': 200,
            'headers': {
                'Content-Type': 'application/json',
                'Access-Control-Allow-Origin': '*'
            },
            'body': json.dumps({
                'predictions': predictions,
                'model_version': get_current_model_version(),
                'inference_time': datetime.now().isoformat()
            })
        }
        
    except Exception as e:
        logger.error(f"Inference failed: {str(e)}")
        return {
            'statusCode': 500,
            'body': json.dumps({'error': str(e)})
        }

# Global model cache to avoid reloading on warm starts
MODEL_CACHE = {}

def load_model_cached():
    """Load model with caching for warm Lambda containers"""
    
    model_version = get_current_model_version()
    
    if model_version in MODEL_CACHE:
        logger.info("Using cached model")
        return MODEL_CACHE[model_version]
    
    # Load model from S3
    model_bucket = os.environ['MODEL_BUCKET']
    model_key = f'models/model-{model_version}.pkl'
    
    logger.info(f"Loading model: s3://{model_bucket}/{model_key}")
    
    response = s3_client.get_object(Bucket=model_bucket, Key=model_key)
    model = pickle.loads(response['Body'].read())
    
    # Cache model
    MODEL_CACHE[model_version] = model
    
    return model

def get_current_model_version() -> str:
    """Get current model version from parameter store"""
    ssm = boto3.client('ssm')
    
    try:
        response = ssm.get_parameter(Name='/ml-model/current-version')
        return response['Parameter']['Value']
    except:
        return 'v1.0.0'  # Default version

def single_inference(model, features: List[float]) -> Dict[str, Any]:
    """Make single prediction"""
    
    input_array = np.array(features).reshape(1, -1)
    prediction = model.predict(input_array)[0]
    confidence = model.predict_proba(input_array).max() if hasattr(model, 'predict_proba') else 0.95
    
    return {
        'prediction': float(prediction),
        'confidence': float(confidence),
        'features_count': len(features)
    }

def batch_inference(model, batch_features: List[List[float]]) -> List[Dict[str, Any]]:
    """Make batch predictions"""
    
    input_array = np.array(batch_features)
    predictions = model.predict(input_array)
    
    if hasattr(model, 'predict_proba'):
        confidences = model.predict_proba(input_array).max(axis=1)
    else:
        confidences = np.ones(len(predictions)) * 0.95
    
    results = []
    for i, (pred, conf) in enumerate(zip(predictions, confidences)):
        results.append({
            'prediction': float(pred),
            'confidence': float(conf),
            'batch_index': i
        })
    
    return results

# ============================================================================
# Lambda Function 4: Model Monitoring & Drift Detection
# ============================================================================
def monitoring_handler(event, context):
    """
    Monitors model performance and detects drift
    Triggered periodically by CloudWatch Events
    """
    try:
        logger.info("Starting model monitoring")
        
        # Get recent predictions and actual outcomes
        monitoring_data = get_monitoring_data()
        
        # Calculate performance metrics
        metrics = calculate_performance_metrics(monitoring_data)
        
        # Detect drift
        drift_metrics = detect_model_drift(monitoring_data)
        
        # Publish metrics to CloudWatch
        publish_monitoring_metrics(metrics, drift_metrics)
        
        # Check if intervention is needed
        if needs_intervention(metrics, drift_metrics):
            trigger_alert_and_retraining(metrics, drift_metrics)
        
        return {
            'statusCode': 200,
            'body': json.dumps({
                'metrics': metrics,
                'drift_metrics': drift_metrics,
                'timestamp': datetime.now().isoformat()
            })
        }
        
    except Exception as e:
        logger.error(f"Monitoring failed: {str(e)}")
        return {
            'statusCode': 500,
            'body': json.dumps({'error': str(e)})
        }

def get_monitoring_data() -> Dict[str, Any]:
    """Retrieve recent predictions and ground truth data"""
    
    # In production, query DynamoDB or other data store
    # For demo, return mock data
    return {
        'predictions': np.random.random(1000).tolist(),
        'actuals': np.random.random(1000).tolist(),
        'timestamps': [datetime.now().isoformat()] * 1000
    }

def calculate_performance_metrics(data: Dict[str, Any]) -> Dict[str, float]:
    """Calculate model performance metrics"""
    
    predictions = np.array(data['predictions'])
    actuals = np.array(data['actuals'])
    
    # Calculate metrics
    mae = np.mean(np.abs(predictions - actuals))
    mse = np.mean((predictions - actuals) ** 2)
    rmse = np.sqrt(mse)
    
    return {
        'mean_absolute_error': float(mae),
        'mean_squared_error': float(mse),
        'root_mean_squared_error': float(rmse),
        'sample_count': len(predictions)
    }

def detect_model_drift(data: Dict[str, Any]) -> Dict[str, float]:
    """Detect statistical drift in predictions"""
    
    predictions = np.array(data['predictions'])
    
    # Compare current distribution with historical baseline
    current_mean = np.mean(predictions)
    current_std = np.std(predictions)
    
    # Historical baseline (in production, load from storage)
    historical_mean = 0.5
    historical_std = 0.2
    
    mean_drift = abs(current_mean - historical_mean) / historical_mean
    std_drift = abs(current_std - historical_std) / historical_std
    
    return {
        'mean_drift': float(mean_drift),
        'std_drift': float(std_drift),
        'drift_score': float(mean_drift + std_drift)
    }

def publish_monitoring_metrics(metrics: Dict[str, float], drift_metrics: Dict[str, float]):
    """Publish metrics to CloudWatch"""
    
    namespace = 'MLOps/ModelPerformance'
    
    # Performance metrics
    for metric_name, value in metrics.items():
        cloudwatch.put_metric_data(
            Namespace=namespace,
            MetricData=[
                {
                    'MetricName': metric_name,
                    'Value': value,
                    'Unit': 'None',
                    'Timestamp': datetime.now()
                }
            ]
        )
    
    # Drift metrics
    for metric_name, value in drift_metrics.items():
        cloudwatch.put_metric_data(
            Namespace=f'{namespace}/Drift',
            MetricData=[
                {
                    'MetricName': metric_name,
                    'Value': value,
                    'Unit': 'None',
                    'Timestamp': datetime.now()
                }
            ]
        )

def needs_intervention(metrics: Dict[str, float], drift_metrics: Dict[str, float]) -> bool:
    """Determine if manual intervention or retraining is needed"""
    
    # Define thresholds
    performance_threshold = 0.1  # Max acceptable RMSE
    drift_threshold = 0.2  # Max acceptable drift score
    
    performance_degraded = metrics.get('root_mean_squared_error', 0) > performance_threshold
    significant_drift = drift_metrics.get('drift_score', 0) > drift_threshold
    
    return performance_degraded or significant_drift

def trigger_alert_and_retraining(metrics: Dict[str, float], drift_metrics: Dict[str, float]):
    """Send alerts and trigger retraining if needed"""
    
    # Send SNS alert
    sns = boto3.client('sns')
    
    message = {
        'alert': 'Model performance degradation detected',
        'metrics': metrics,
        'drift_metrics': drift_metrics,
        'timestamp': datetime.now().isoformat()
    }
    
    sns.publish(
        TopicArn=os.environ['ALERT_TOPIC_ARN'],
        Message=json.dumps(message),
        Subject='MLOps Alert: Model Intervention Required'
    )
    
    # Trigger retraining workflow
    if drift_metrics.get('drift_score', 0) > 0.3:  # High drift threshold
        stepfunctions.start_execution(
            stateMachineArn=os.environ['RETRAINING_WORKFLOW_ARN'],
            name=f"emergency-retrain-{datetime.now().strftime('%Y%m%d-%H%M%S')}",
            input=json.dumps({
                'trigger': 'drift_detection',
                'metrics': metrics,
                'drift_metrics': drift_metrics
            })
        )

def publish_data_metrics(processed_data: Dict[str, Any]):
    """Publish data processing metrics to CloudWatch"""
    
    metrics = [
        {
            'MetricName': 'DataQualityScore',
            'Value': processed_data['quality_score'],
            'Unit': 'Percent'
        },
        {
            'MetricName': 'FeatureDrift',
            'Value': processed_data['feature_drift'],
            'Unit': 'None'
        },
        {
            'MetricName': 'ProcessedRows',
            'Value': processed_data['processed_rows'],
            'Unit': 'Count'
        }
    ]
    
    cloudwatch.put_metric_data(
        Namespace='MLOps/DataProcessing',
        MetricData=[{
            **metric,
            'Timestamp': datetime.now()
        } for metric in metrics]
    )

def log_inference_metrics(prediction_count: int):
    """Log inference metrics"""
    
    cloudwatch.put_metric_data(
        Namespace='MLOps/Inference',
        MetricData=[
            {
                'MetricName': 'PredictionCount',
                'Value': prediction_count,
                'Unit': 'Count',
                'Timestamp': datetime.now()
            },
            {
                'MetricName': 'InvocationCount',
                'Value': 1,
                'Unit': 'Count',
                'Timestamp': datetime.now()
            }
        ]
    )

# ============================================================================
# Infrastructure as Code (Serverless Framework / CDK)
# ============================================================================

# Example serverless.yml configuration
SERVERLESS_CONFIG = '''
service: serverless-mlops

provider:
  name: aws
  runtime: python3.9
  region: us-east-1
  environment:
    MODEL_BUCKET: ${self:service}-models-${sls:stage}
    TRAINING_WORKFLOW_ARN: !Ref TrainingStateMachine
    ALERT_TOPIC_ARN: !Ref AlertTopic
    
functions:
  dataProcessing:
    handler: handler.data_processing_handler
    timeout: 900
    memorySize: 1024
    events:
      - s3:
          bucket: ${self:provider.environment.MODEL_BUCKET}
          event: s3:ObjectCreated:*
          rules:
            - prefix: raw-data/
            - suffix: .csv
            
  trainingOrchestrator:
    handler: handler.training_orchestrator_handler
    timeout: 300
    memorySize: 512
    environment:
      TRAINING_WORKFLOW_ARN: !Ref TrainingStateMachine
      
  inference:
    handler: handler.inference_handler
    timeout: 30
    memorySize: 2048
    provisionedConcurrency: 5  # Avoid cold starts for inference
    events:
      - http:
          path: /predict
          method: post
          cors: true
      - http:
          path: /batch-predict  
          method: post
          cors: true
          
  monitoring:
    handler: handler.monitoring_handler
    timeout: 600
    memorySize: 1024
    events:
      - schedule: rate(1 hour)  # Run monitoring every hour

resources:
  Resources:
    AlertTopic:
      Type: AWS::SNS::Topic
      Properties:
        DisplayName: MLOps Alerts
        
    TrainingStateMachine:
      Type: AWS::StepFunctions::StateMachine
      Properties:
        StateMachineName: ${self:service}-training-${sls:stage}
        RoleArn: !GetAtt StepFunctionsRole.Arn
        DefinitionString: |
          {
            "Comment": "ML Training Pipeline",
            "StartAt": "StartTraining",
            "States": {
              "StartTraining": {
                "Type": "Task",
                "Resource": "arn:aws:states:::sagemaker:createTrainingJob.sync",
                "Parameters": {
                  "TrainingJobName.$": "$.training_job_name",
                  "AlgorithmSpecification.$": "$.algorithm_specification",
                  "InputDataConfig.$": "$.input_data_config",
                  "OutputDataConfig.$": "$.output_data_config",
                  "ResourceConfig.$": "$.resource_config",
                  "StoppingCondition.$": "$.stopping_condition",
                  "RoleArn.$": "$.role_arn"
                },
                "Next": "DeployModel"
              },
              "DeployModel": {
                "Type": "Task", 
                "Resource": "arn:aws:states:::lambda:invoke",
                "Parameters": {
                  "FunctionName": "${self:service}-${sls:stage}-modelDeployment"
                },
                "End": true
              }
            }
          }
'''

print("Serverless MLOps system configured for production deployment")