Securing AWS AI/ML Pipelines: DevSecOps Integration Guide

Article Content

• Article Content
• Introduction: The Critical Need for MLOps Security
• Understanding ML Pipeline Security Threats
  • The ML Attack Surface
  • AWS ML Security Threat Model
• Secure Data Pipeline Implementation
  • Data Security Throughout the ML Lifecycle
    • Comprehensive Data Protection Framework
• Secure Model Training and Development
  • Training Environment Security
    • Secure SageMaker Training Implementation
• Model Deployment Security
  • Secure Inference Endpoints
    • Secure Endpoint Deployment
• Implementation Roadmap for Secure ML Pipelines
  • Phase 1: Data Security Foundation (Weeks 1-3)
  • Phase 2: Training Security (Weeks 4-6)
  • Phase 3: Deployment Security (Weeks 7-9)
  • Phase 4: Monitoring and Governance (Weeks 10-12)
• Related Articles and Additional Resources
  • AWS Documentation
  • Industry Standards
  • Community Resources

Introduction: The Critical Need for MLOps Security

The integration of machine learning into production systems has revolutionized enterprise applications, with 87% of organizations now running ML models in production environments. However, this rapid adoption has created significant security gaps: 71% of ML pipelines lack adequate security controls, and 45% of data scientists report having insufficient security training for production deployments.

Traditional DevSecOps practices don’t directly translate to ML workflows due to unique challenges including sensitive training data, model intellectual property, and complex multi-stage pipelines. AWS provides comprehensive services for MLOps, but securing these pipelines requires specialized knowledge of both security best practices and ML-specific attack vectors.

Recent security incidents highlight the urgency: 34% of ML attacks target the training pipeline, 28% focus on data poisoning, and 23% attempt model extraction. The financial impact averages $4.1 million per incident, making MLOps security a critical business imperative.

This guide provides practical, tested implementations for securing AI/ML pipelines on AWS throughout the entire development lifecycle. We’ll cover data security, model protection, deployment safety, and monitoring strategies with working code examples and enterprise-grade configurations.

Understanding ML Pipeline Security Threats

The ML Attack Surface

AI/ML pipelines introduce unique security challenges that extend beyond traditional application security:

Data Pipeline Vulnerabilities: Training datasets often contain sensitive information and represent valuable intellectual property. 67% of ML security incidents involve data exfiltration or manipulation during the data preparation phase.

Model Intellectual Property: Trained models represent significant business value and competitive advantage. Model extraction attacks can steal proprietary algorithms with 89% success rates against unprotected endpoints.

Training Infrastructure: ML training requires significant computational resources, making infrastructure a target for both resource theft and supply chain attacks. 52% of container-based ML deployments contain vulnerabilities.

Deployment Attack Vectors: ML models in production face adversarial attacks, model poisoning, and inference manipulation. Real-time inference endpoints experience 2.3x more attack attempts than traditional web applications.

AWS ML Security Threat Model

import boto3
import json
from datetime import datetime, timedelta
from typing import Dict, List, Optional, Tuple
import hashlib
import os

class MLPipelineThreatModel:
    """
    Comprehensive threat modeling for AWS ML pipelines
    Identifies and categorizes security risks throughout the ML lifecycle
    """
    
    def __init__(self, region_name: str = 'us-east-1'):
        self.sagemaker = boto3.client('sagemaker', region_name=region_name)
        self.s3 = boto3.client('s3', region_name=region_name)
        self.iam = boto3.client('iam', region_name=region_name)
        self.kms = boto3.client('kms', region_name=region_name)
        
        # Initialize threat categories
        self.threat_categories = {
            'data_security': [
                'unauthorized_data_access',
                'data_poisoning',
                'data_exfiltration',
                'privacy_violations'
            ],
            'model_security': [
                'model_theft',
                'model_poisoning',
                'adversarial_attacks',
                'model_inversion'
            ],
            'infrastructure_security': [
                'container_vulnerabilities',
                'privilege_escalation',
                'resource_hijacking',
                'supply_chain_attacks'
            ],
            'deployment_security': [
                'inference_manipulation',
                'endpoint_exploitation',
                'monitoring_bypass',
                'rollback_attacks'
            ]
        }
    
    def assess_pipeline_threats(self, pipeline_config: Dict) -> Dict:
        """Comprehensive threat assessment for ML pipeline"""
        
        assessment_results = {
            'pipeline_id': pipeline_config.get('pipeline_name', 'unknown'),
            'assessment_timestamp': datetime.utcnow().isoformat(),
            'overall_risk_score': 0,
            'threat_analysis': {},
            'recommendations': [],
            'compliance_gaps': []
        }
        
        # Analyze each component of the ML pipeline
        for component, config in pipeline_config.items():
            if component == 'data_processing':
                data_threats = self._assess_data_security_threats(config)
                assessment_results['threat_analysis']['data_security'] = data_threats
                
            elif component == 'training':
                training_threats = self._assess_training_security_threats(config)
                assessment_results['threat_analysis']['model_security'] = training_threats
                
            elif component == 'deployment':
                deployment_threats = self._assess_deployment_security_threats(config)
                assessment_results['threat_analysis']['deployment_security'] = deployment_threats
                
            elif component == 'infrastructure':
                infra_threats = self._assess_infrastructure_security_threats(config)
                assessment_results['threat_analysis']['infrastructure_security'] = infra_threats
        
        # Calculate overall risk score
        assessment_results['overall_risk_score'] = self._calculate_overall_risk_score(
            assessment_results['threat_analysis']
        )
        
        # Generate recommendations
        assessment_results['recommendations'] = self._generate_security_recommendations(
            assessment_results['threat_analysis']
        )
        
        return assessment_results
    
    def _assess_data_security_threats(self, data_config: Dict) -> Dict:
        """Assess threats specific to data processing components"""
        
        threats = {
            'risk_level': 'low',
            'identified_threats': [],
            'security_gaps': [],
            'recommendations': []
        }
        
        # Check data encryption
        if not data_config.get('encryption_enabled', False):
            threats['identified_threats'].append({
                'threat_type': 'data_exposure',
                'severity': 'high',
                'description': 'Data stored without encryption',
                'impact': 'Unauthorized access to sensitive training data'
            })
            threats['risk_level'] = 'high'
        
        # Check access controls
        if not data_config.get('access_controls', {}):
            threats['identified_threats'].append({
                'threat_type': 'unauthorized_access',
                'severity': 'medium',
                'description': 'Insufficient access controls on data',
                'impact': 'Potential data manipulation or theft'
            })
        
        # Check data lineage tracking
        if not data_config.get('lineage_tracking', False):
            threats['security_gaps'].append({
                'gap_type': 'audit_trail',
                'description': 'No data lineage tracking implemented',
                'recommendation': 'Implement comprehensive data lineage tracking'
            })
        
        # Check for PII handling
        if data_config.get('contains_pii', False) and not data_config.get('pii_protection', False):
            threats['identified_threats'].append({
                'threat_type': 'privacy_violation',
                'severity': 'critical',
                'description': 'PII data without adequate protection',
                'impact': 'Regulatory compliance violations and privacy breaches'
            })
            threats['risk_level'] = 'critical'
        
        return threats
    
    def _assess_training_security_threats(self, training_config: Dict) -> Dict:
        """Assess threats specific to model training components"""
        
        threats = {
            'risk_level': 'low',
            'identified_threats': [],
            'security_gaps': [],
            'recommendations': []
        }
        
        # Check training job isolation
        if not training_config.get('network_isolation', False):
            threats['identified_threats'].append({
                'threat_type': 'training_exposure',
                'severity': 'medium',
                'description': 'Training jobs not network isolated',
                'impact': 'Potential model theft or manipulation'
            })
        
        # Check model artifact protection
        if not training_config.get('model_encryption', False):
            threats['identified_threats'].append({
                'threat_type': 'model_theft',
                'severity': 'high',
                'description': 'Model artifacts not encrypted',
                'impact': 'Intellectual property theft'
            })
            threats['risk_level'] = 'high'
        
        # Check training data validation
        if not training_config.get('data_validation', False):
            threats['identified_threats'].append({
                'threat_type': 'data_poisoning',
                'severity': 'high',
                'description': 'No training data validation implemented',
                'impact': 'Model performance degradation or backdoors'
            })
        
        # Check hyperparameter security
        if training_config.get('exposed_hyperparameters', False):
            threats['security_gaps'].append({
                'gap_type': 'information_leakage',
                'description': 'Hyperparameters exposed in logs or metadata',
                'recommendation': 'Sanitize hyperparameter logging'
            })
        
        return threats
    
    def _assess_deployment_security_threats(self, deployment_config: Dict) -> Dict:
        """Assess threats specific to model deployment components"""
        
        threats = {
            'risk_level': 'low',
            'identified_threats': [],
            'security_gaps': [],
            'recommendations': []
        }
        
        # Check endpoint security
        if not deployment_config.get('authentication_enabled', False):
            threats['identified_threats'].append({
                'threat_type': 'unauthorized_inference',
                'severity': 'high',
                'description': 'Model endpoints without authentication',
                'impact': 'Unauthorized model access and potential data extraction'
            })
            threats['risk_level'] = 'high'
        
        # Check inference monitoring
        if not deployment_config.get('inference_monitoring', False):
            threats['security_gaps'].append({
                'gap_type': 'attack_detection',
                'description': 'No inference request monitoring',
                'recommendation': 'Implement real-time inference monitoring'
            })
        
        # Check model versioning
        if not deployment_config.get('version_control', False):
            threats['identified_threats'].append({
                'threat_type': 'rollback_attack',
                'severity': 'medium',
                'description': 'No model version control',
                'impact': 'Potential deployment of compromised models'
            })
        
        # Check rate limiting
        if not deployment_config.get('rate_limiting', False):
            threats['identified_threats'].append({
                'threat_type': 'resource_exhaustion',
                'severity': 'medium',
                'description': 'No rate limiting on inference endpoints',
                'impact': 'Denial of service attacks'
            })
        
        return threats
    
    def _assess_infrastructure_security_threats(self, infra_config: Dict) -> Dict:
        """Assess threats specific to infrastructure components"""
        
        threats = {
            'risk_level': 'low',
            'identified_threats': [],
            'security_gaps': [],
            'recommendations': []
        }
        
        # Check container security
        if infra_config.get('uses_containers', False):
            if not infra_config.get('container_scanning', False):
                threats['identified_threats'].append({
                    'threat_type': 'container_vulnerability',
                    'severity': 'high',
                    'description': 'Containers not scanned for vulnerabilities',
                    'impact': 'Potential exploitation of container vulnerabilities'
                })
        
        # Check IAM permissions
        if infra_config.get('overprivileged_roles', False):
            threats['identified_threats'].append({
                'threat_type': 'privilege_escalation',
                'severity': 'high',
                'description': 'Overprivileged IAM roles detected',
                'impact': 'Potential lateral movement and privilege escalation'
            })
            threats['risk_level'] = 'high'
        
        # Check network security
        if not infra_config.get('vpc_isolation', False):
            threats['security_gaps'].append({
                'gap_type': 'network_isolation',
                'description': 'Infrastructure not isolated in VPC',
                'recommendation': 'Deploy ML infrastructure in isolated VPC'
            })
        
        return threats
    
    def _calculate_overall_risk_score(self, threat_analysis: Dict) -> int:
        """Calculate overall risk score based on identified threats"""
        
        risk_scores = {
            'critical': 10,
            'high': 7,
            'medium': 4,
            'low': 1
        }
        
        total_score = 0
        threat_count = 0
        
        for category, threats in threat_analysis.items():
            if isinstance(threats, dict) and 'identified_threats' in threats:
                for threat in threats['identified_threats']:
                    severity = threat.get('severity', 'low')
                    total_score += risk_scores.get(severity, 1)
                    threat_count += 1
        
        if threat_count == 0:
            return 1
        
        average_score = total_score / threat_count
        return min(10, max(1, int(average_score)))
    
    def _generate_security_recommendations(self, threat_analysis: Dict) -> List[Dict]:
        """Generate prioritized security recommendations"""
        
        recommendations = []
        
        # High-priority recommendations based on critical/high threats
        for category, threats in threat_analysis.items():
            if isinstance(threats, dict) and 'identified_threats' in threats:
                for threat in threats['identified_threats']:
                    if threat.get('severity') in ['critical', 'high']:
                        recommendations.append({
                            'priority': 'high',
                            'category': category,
                            'threat_type': threat['threat_type'],
                            'recommendation': self._get_threat_remediation(threat['threat_type']),
                            'estimated_effort': self._estimate_remediation_effort(threat['threat_type'])
                        })
        
        return recommendations
    
    def _get_threat_remediation(self, threat_type: str) -> str:
        """Get specific remediation advice for threat types"""
        
        remediations = {
            'data_exposure': 'Implement KMS encryption for all data at rest and in transit',
            'unauthorized_access': 'Deploy IAM policies with least privilege access',
            'data_poisoning': 'Implement data validation and integrity checks',
            'model_theft': 'Enable model artifact encryption and access logging',
            'training_exposure': 'Enable network isolation for training jobs',
            'unauthorized_inference': 'Implement authentication and authorization for endpoints',
            'container_vulnerability': 'Enable container image scanning and vulnerability patching',
            'privilege_escalation': 'Review and restrict IAM permissions to minimum required',
            'privacy_violation': 'Implement PII detection and anonymization',
            'resource_exhaustion': 'Deploy rate limiting and resource quotas'
        }
        
        return remediations.get(threat_type, 'Consult security team for specific guidance')
    
    def _estimate_remediation_effort(self, threat_type: str) -> str:
        """Estimate effort required for threat remediation"""
        
        effort_estimates = {
            'data_exposure': 'medium',
            'unauthorized_access': 'high',
            'data_poisoning': 'high',
            'model_theft': 'medium',
            'training_exposure': 'low',
            'unauthorized_inference': 'medium',
            'container_vulnerability': 'medium',
            'privilege_escalation': 'high',
            'privacy_violation': 'high',
            'resource_exhaustion': 'low'
        }
        
        return effort_estimates.get(threat_type, 'unknown')
    
    def generate_security_report(self, assessment_results: Dict) -> str:
        """Generate comprehensive security assessment report"""
        
        report = f"""
# ML Pipeline Security Assessment Report

**Pipeline ID**: {assessment_results['pipeline_id']}
**Assessment Date**: {assessment_results['assessment_timestamp']}
**Overall Risk Score**: {assessment_results['overall_risk_score']}/10

## Executive Summary

This assessment identified {sum(len(threats.get('identified_threats', [])) for threats in assessment_results['threat_analysis'].values())} security threats across the ML pipeline.

## Threat Analysis by Category

"""
        
        for category, threats in assessment_results['threat_analysis'].items():
            report += f"""
### {category.replace('_', ' ').title()}

**Risk Level**: {threats.get('risk_level', 'unknown')}
**Identified Threats**: {len(threats.get('identified_threats', []))}

"""
            
            for threat in threats.get('identified_threats', []):
                report += f"""
- **{threat['threat_type']}** ({threat['severity']}): {threat['description']}
  - Impact: {threat['impact']}

"""
        
        report += """
## Priority Recommendations

"""
        
        high_priority = [r for r in assessment_results['recommendations'] if r['priority'] == 'high']
        for i, rec in enumerate(high_priority[:5], 1):
            report += f"""
{i}. **{rec['category']}**: {rec['recommendation']}
   - Effort: {rec['estimated_effort']}
   - Addresses: {rec['threat_type']}

"""
        
        return report

Secure Data Pipeline Implementation

Data Security Throughout the ML Lifecycle

Securing data is fundamental to ML pipeline security, requiring protection from collection through model training and inference.

Comprehensive Data Protection Framework

import boto3
import json
import pandas as pd
import numpy as np
from datetime import datetime, timedelta
from typing import Dict, List, Optional, Any
import hashlib
import os
import re

class SecureDataPipeline:
    """
    Secure data pipeline implementation for ML workloads
    Handles data encryption, validation, lineage tracking, and privacy protection
    """
    
    def __init__(self, region_name: str = 'us-east-1'):
        self.s3 = boto3.client('s3', region_name=region_name)
        self.kms = boto3.client('kms', region_name=region_name)
        self.glue = boto3.client('glue', region_name=region_name)
        self.sagemaker = boto3.client('sagemaker', region_name=region_name)
        self.logs = boto3.client('logs', region_name=region_name)
        
        # Data classification levels
        self.data_classifications = {
            'public': {'encryption_required': False, 'access_level': 'open'},
            'internal': {'encryption_required': True, 'access_level': 'restricted'},
            'confidential': {'encryption_required': True, 'access_level': 'limited'},
            'restricted': {'encryption_required': True, 'access_level': 'minimal'}
        }
    
    def create_secure_data_bucket(self, 
                                  bucket_name: str,
                                  kms_key_id: str,
                                  data_classification: str = 'confidential') -> Dict:
        """Create secure S3 bucket for ML data with comprehensive protection"""
        
        classification_config = self.data_classifications.get(data_classification, 
                                                              self.data_classifications['confidential'])
        
        # Create bucket with security configurations
        try:
            # Create the bucket
            if boto3.Session().region_name != 'us-east-1':
                self.s3.create_bucket(
                    Bucket=bucket_name,
                    CreateBucketConfiguration={'LocationConstraint': boto3.Session().region_name}
                )
            else:
                self.s3.create_bucket(Bucket=bucket_name)
            
            # Enable versioning
            self.s3.put_bucket_versioning(
                Bucket=bucket_name,
                VersioningConfiguration={'Status': 'Enabled'}
            )
            
            # Configure encryption
            if classification_config['encryption_required']:
                self.s3.put_bucket_encryption(
                    Bucket=bucket_name,
                    ServerSideEncryptionConfiguration={
                        'Rules': [
                            {
                                'ApplyServerSideEncryptionByDefault': {
                                    'SSEAlgorithm': 'aws:kms',
                                    'KMSMasterKeyID': kms_key_id
                                },
                                'BucketKeyEnabled': True
                            }
                        ]
                    }
                )
            
            # Block public access
            self.s3.put_public_access_block(
                Bucket=bucket_name,
                PublicAccessBlockConfiguration={
                    'BlockPublicAcls': True,
                    'IgnorePublicAcls': True,
                    'BlockPublicPolicy': True,
                    'RestrictPublicBuckets': True
                }
            )
            
            # Enable logging
            self.s3.put_bucket_logging(
                Bucket=bucket_name,
                BucketLoggingStatus={
                    'LoggingEnabled': {
                        'TargetBucket': f'{bucket_name}-access-logs',
                        'TargetPrefix': 'access-logs/'
                    }
                }
            )
            
            # Configure lifecycle policies
            self.s3.put_bucket_lifecycle_configuration(
                Bucket=bucket_name,
                LifecycleConfiguration={
                    'Rules': [
                        {
                            'ID': 'ml-data-lifecycle',
                            'Status': 'Enabled',
                            'Filter': {'Prefix': 'training-data/'},
                            'Transitions': [
                                {
                                    'Days': 30,
                                    'StorageClass': 'STANDARD_IA'
                                },
                                {
                                    'Days': 90,
                                    'StorageClass': 'GLACIER'
                                }
                            ]
                        }
                    ]
                }
            )
            
            # Set bucket policy based on classification
            bucket_policy = self._generate_bucket_policy(bucket_name, data_classification)
            self.s3.put_bucket_policy(
                Bucket=bucket_name,
                Policy=json.dumps(bucket_policy)
            )
            
            # Enable notifications for security monitoring
            self.s3.put_bucket_notification_configuration(
                Bucket=bucket_name,
                NotificationConfiguration={
                    'CloudWatchConfigurations': [
                        {
                            'Id': 'security-monitoring',
                            'CloudWatchConfiguration': {
                                'LogGroupName': f'/aws/s3/{bucket_name}',
                                'Event': 's3:ObjectCreated:*'
                            },
                            'Filter': {
                                'Key': {
                                    'FilterRules': [
                                        {
                                            'Name': 'prefix',
                                            'Value': 'sensitive-data/'
                                        }
                                    ]
                                }
                            }
                        }
                    ]
                }
            )
            
            return {
                'bucket_name': bucket_name,
                'encryption_enabled': classification_config['encryption_required'],
                'kms_key_id': kms_key_id,
                'classification': data_classification,
                'security_features': [
                    'versioning_enabled',
                    'public_access_blocked',
                    'access_logging_enabled',
                    'lifecycle_policies_configured',
                    'security_monitoring_enabled'
                ]
            }
            
        except Exception as e:
            print(f"Error creating secure bucket: {e}")
            raise
    
    def _generate_bucket_policy(self, bucket_name: str, classification: str) -> Dict:
        """Generate bucket policy based on data classification"""
        
        base_policy = {
            "Version": "2012-10-17",
            "Statement": [
                {
                    "Sid": "DenyInsecureConnections",
                    "Effect": "Deny",
                    "Principal": "*",
                    "Action": "s3:*",
                    "Resource": [
                        f"arn:aws:s3:::{bucket_name}/*",
                        f"arn:aws:s3:::{bucket_name}"
                    ],
                    "Condition": {
                        "Bool": {
                            "aws:SecureTransport": "false"
                        }
                    }
                }
            ]
        }
        
        if classification in ['confidential', 'restricted']:
            # Add IP restriction for highly sensitive data
            base_policy["Statement"].append({
                "Sid": "RestrictByIP",
                "Effect": "Deny",
                "Principal": "*",
                "Action": "s3:*",
                "Resource": [
                    f"arn:aws:s3:::{bucket_name}/*",
                    f"arn:aws:s3:::{bucket_name}"
                ],
                "Condition": {
                    "NotIpAddress": {
                        "aws:SourceIp": [
                            "10.0.0.0/8",
                            "172.16.0.0/12",
                            "192.168.0.0/16"
                        ]
                    }
                }
            })
        
        if classification == 'restricted':
            # Add MFA requirement for restricted data
            base_policy["Statement"].append({
                "Sid": "RequireMFA",
                "Effect": "Deny",
                "Principal": "*",
                "Action": "s3:*",
                "Resource": [
                    f"arn:aws:s3:::{bucket_name}/*",
                    f"arn:aws:s3:::{bucket_name}"
                ],
                "Condition": {
                    "BoolIfExists": {
                        "aws:MultiFactorAuthPresent": "false"
                    }
                }
            })
        
        return base_policy
    
    def implement_data_validation(self, 
                                  data_source: str,
                                  validation_rules: Dict,
                                  output_location: str) -> Dict:
        """Implement comprehensive data validation for ML pipelines"""
        
        validation_job_name = f"data-validation-{datetime.utcnow().strftime('%Y%m%d-%H%M%S')}"
        
        # Create Glue job for data validation
        job_definition = {
            'Name': validation_job_name,
            'Role': 'arn:aws:iam::123456789012:role/GlueDataValidationRole',
            'Command': {
                'Name': 'glueetl',
                'ScriptLocation': 's3://ml-security-scripts/data_validation.py',
                'PythonVersion': '3'
            },
            'DefaultArguments': {
                '--job-language': 'python',
                '--data_source': data_source,
                '--validation_rules': json.dumps(validation_rules),
                '--output_location': output_location,
                '--enable-metrics': '',
                '--enable-continuous-cloudwatch-log': 'true',
                '--enable-spark-ui': 'true',
                '--spark-event-logs-path': f's3://ml-security-logs/spark-events/{validation_job_name}/'
            },
            'MaxRetries': 1,
            'Timeout': 2880,  # 48 hours
            'MaxCapacity': 10.0,
            'SecurityConfiguration': 'ml-security-configuration',
            'Tags': {
                'Purpose': 'DataValidation',
                'Environment': 'production',
                'DataClassification': validation_rules.get('classification', 'confidential')
            }
        }
        
        try:
            response = self.glue.create_job(**job_definition)
            
            # Start the validation job
            job_run_response = self.glue.start_job_run(
                JobName=validation_job_name,
                Arguments={
                    '--validation_timestamp': datetime.utcnow().isoformat(),
                    '--security_scan_enabled': 'true'
                }
            )
            
            return {
                'job_name': validation_job_name,
                'job_run_id': job_run_response['JobRunId'],
                'validation_rules': validation_rules,
                'status': 'started'
            }
            
        except Exception as e:
            print(f"Error creating validation job: {e}")
            raise
    
    def create_data_validation_script(self) -> str:
        """Generate comprehensive data validation script for Glue"""
        
        validation_script = '''
import sys
import boto3
import pandas as pd
import numpy as np
import json
import re
from datetime import datetime
from awsglue.transforms import *
from awsglue.utils import getResolvedOptions
from pyspark.context import SparkContext
from awsglue.context import GlueContext
from awsglue.job import Job

# Initialize Glue context
args = getResolvedOptions(sys.argv, ['JOB_NAME', 'data_source', 'validation_rules', 'output_location'])
sc = SparkContext()
glueContext = GlueContext(sc)
spark = glueContext.spark_session
job = Job(glueContext)
job.init(args['JOB_NAME'], args)

class MLDataValidator:
    def __init__(self, spark_session):
        self.spark = spark_session
        self.validation_results = {
            'timestamp': datetime.utcnow().isoformat(),
            'total_records': 0,
            'validation_passed': True,
            'errors': [],
            'warnings': [],
            'security_issues': [],
            'quality_metrics': {}
        }
    
    def validate_data_quality(self, df, rules):
        """Comprehensive data quality validation"""
        
        self.validation_results['total_records'] = df.count()
        
        # Schema validation
        self._validate_schema(df, rules.get('schema_rules', {}))
        
        # Data type validation
        self._validate_data_types(df, rules.get('type_rules', {}))
        
        # Range validation
        self._validate_ranges(df, rules.get('range_rules', {}))
        
        # Null value validation
        self._validate_null_values(df, rules.get('null_rules', {}))
        
        # Duplicate detection
        self._detect_duplicates(df, rules.get('duplicate_rules', {}))
        
        # Security validation
        self._validate_security(df, rules.get('security_rules', {}))
        
        # Statistical validation
        self._validate_statistics(df, rules.get('statistical_rules', {}))
        
        return self.validation_results
    
    def _validate_schema(self, df, schema_rules):
        """Validate DataFrame schema against expected schema"""
        
        expected_columns = schema_rules.get('required_columns', [])
        actual_columns = df.columns
        
        missing_columns = set(expected_columns) - set(actual_columns)
        if missing_columns:
            self.validation_results['errors'].append({
                'type': 'schema_validation',
                'message': f'Missing required columns: {missing_columns}',
                'severity': 'high'
            })
            self.validation_results['validation_passed'] = False
        
        # Check for unexpected columns
        allowed_columns = schema_rules.get('allowed_columns', actual_columns)
        unexpected_columns = set(actual_columns) - set(allowed_columns)
        if unexpected_columns:
            self.validation_results['warnings'].append({
                'type': 'schema_validation',
                'message': f'Unexpected columns found: {unexpected_columns}',
                'severity': 'medium'
            })
    
    def _validate_data_types(self, df, type_rules):
        """Validate data types for each column"""
        
        for column, expected_type in type_rules.items():
            if column in df.columns:
                actual_type = str(df.schema[column].dataType)
                if expected_type.lower() not in actual_type.lower():
                    self.validation_results['errors'].append({
                        'type': 'data_type_validation',
                        'message': f'Column {column} has type {actual_type}, expected {expected_type}',
                        'severity': 'high'
                    })
                    self.validation_results['validation_passed'] = False
    
    def _validate_ranges(self, df, range_rules):
        """Validate numerical ranges for specified columns"""
        
        for column, range_config in range_rules.items():
            if column in df.columns:
                min_val = range_config.get('min')
                max_val = range_config.get('max')
                
                if min_val is not None:
                    out_of_range_count = df.filter(df[column] < min_val).count()
                    if out_of_range_count > 0:
                        self.validation_results['errors'].append({
                            'type': 'range_validation',
                            'message': f'Column {column}: {out_of_range_count} values below minimum {min_val}',
                            'severity': 'medium'
                        })
                
                if max_val is not None:
                    out_of_range_count = df.filter(df[column] > max_val).count()
                    if out_of_range_count > 0:
                        self.validation_results['errors'].append({
                            'type': 'range_validation',
                            'message': f'Column {column}: {out_of_range_count} values above maximum {max_val}',
                            'severity': 'medium'
                        })
    
    def _validate_null_values(self, df, null_rules):
        """Validate null value constraints"""
        
        for column, null_config in null_rules.items():
            if column in df.columns:
                null_count = df.filter(df[column].isNull()).count()
                
                if not null_config.get('allowed', True) and null_count > 0:
                    self.validation_results['errors'].append({
                        'type': 'null_validation',
                        'message': f'Column {column} contains {null_count} null values (not allowed)',
                        'severity': 'high'
                    })
                    self.validation_results['validation_passed'] = False
                
                max_null_percentage = null_config.get('max_percentage', 100)
                null_percentage = (null_count / self.validation_results['total_records']) * 100
                
                if null_percentage > max_null_percentage:
                    self.validation_results['warnings'].append({
                        'type': 'null_validation',
                        'message': f'Column {column} has {null_percentage:.2f}% null values, exceeds {max_null_percentage}%',
                        'severity': 'medium'
                    })
    
    def _detect_duplicates(self, df, duplicate_rules):
        """Detect and report duplicate records"""
        
        key_columns = duplicate_rules.get('key_columns', [])
        if key_columns:
            original_count = df.count()
            unique_count = df.dropDuplicates(key_columns).count()
            duplicate_count = original_count - unique_count
            
            if duplicate_count > 0:
                max_duplicates = duplicate_rules.get('max_duplicates', 0)
                severity = 'high' if duplicate_count > max_duplicates else 'medium'
                
                self.validation_results['warnings'].append({
                    'type': 'duplicate_detection',
                    'message': f'Found {duplicate_count} duplicate records based on {key_columns}',
                    'severity': severity
                })
                
                self.validation_results['quality_metrics']['duplicate_percentage'] = (duplicate_count / original_count) * 100
    
    def _validate_security(self, df, security_rules):
        """Validate security constraints and detect sensitive data"""
        
        # PII detection patterns
        pii_patterns = {
            'ssn': r'\\b\\d{3}-\\d{2}-\\d{4}\\b',
            'credit_card': r'\\b\\d{4}[- ]?\\d{4}[- ]?\\d{4}[- ]?\\d{4}\\b',
            'email': r'\\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\\.[A-Z|a-z]{2,}\\b',
            'phone': r'\\b\\d{3}[- ]?\\d{3}[- ]?\\d{4}\\b'
        }
        
        text_columns = [col for col, dtype in df.dtypes if dtype == 'string']
        
        for column in text_columns:
            if security_rules.get('scan_for_pii', True):
                for pii_type, pattern in pii_patterns.items():
                    # Sample data for PII detection (avoid scanning entire dataset)
                    sample_df = df.sample(fraction=0.1).select(column).collect()
                    
                    pii_found = False
                    for row in sample_df[:1000]:  # Limit sample size
                        if row[column] and re.search(pattern, str(row[column])):
                            pii_found = True
                            break
                    
                    if pii_found:
                        self.validation_results['security_issues'].append({
                            'type': 'pii_detection',
                            'message': f'Potential {pii_type} detected in column {column}',
                            'severity': 'critical'
                        })
                        self.validation_results['validation_passed'] = False
        
        # Check for sensitive column names
        sensitive_keywords = ['password', 'secret', 'key', 'token', 'credential']
        for column in df.columns:
            if any(keyword in column.lower() for keyword in sensitive_keywords):
                self.validation_results['security_issues'].append({
                    'type': 'sensitive_column',
                    'message': f'Column name {column} suggests sensitive data',
                    'severity': 'high'
                })
    
    def _validate_statistics(self, df, statistical_rules):
        """Validate statistical properties of the data"""
        
        numeric_columns = [col for col, dtype in df.dtypes if dtype in ['int', 'double', 'float']]
        
        for column in numeric_columns:
            if column in statistical_rules:
                rules = statistical_rules[column]
                
                # Calculate statistics
                stats = df.select(column).describe().collect()
                stats_dict = {row['summary']: float(row[column]) for row in stats if row[column] != column}
                
                # Validate mean
                if 'expected_mean' in rules:
                    expected_mean = rules['expected_mean']
                    tolerance = rules.get('mean_tolerance', 0.1)
                    actual_mean = stats_dict.get('mean', 0)
                    
                    if abs(actual_mean - expected_mean) > expected_mean * tolerance:
                        self.validation_results['warnings'].append({
                            'type': 'statistical_validation',
                            'message': f'Column {column} mean {actual_mean:.2f} differs from expected {expected_mean:.2f}',
                            'severity': 'medium'
                        })
                
                # Store quality metrics
                self.validation_results['quality_metrics'][f'{column}_statistics'] = stats_dict

# Main execution
try:
    # Load validation rules
    validation_rules = json.loads(args['validation_rules'])
    
    # Read data
    datasource = glueContext.create_dynamic_frame.from_options(
        's3',
        {'paths': [args['data_source']]},
        format='json'
    )
    
    df = datasource.toDF()
    
    # Initialize validator
    validator = MLDataValidator(spark)
    
    # Run validation
    results = validator.validate_data_quality(df, validation_rules)
    
    # Save validation results
    results_df = spark.createDataFrame([results])
    results_output = glueContext.create_dynamic_frame.from_dataframe(results_df, glueContext, "validation_results")
    
    glueContext.write_dynamic_frame.from_options(
        frame=results_output,
        connection_type="s3",
        connection_options={"path": f"{args['output_location']}/validation_results/"},
        format="json"
    )
    
    # If validation passed, write cleaned data
    if results['validation_passed']:
        cleaned_output = glueContext.create_dynamic_frame.from_dataframe(df, glueContext, "cleaned_data")
        glueContext.write_dynamic_frame.from_options(
            frame=cleaned_output,
            connection_type="s3",
            connection_options={"path": f"{args['output_location']}/validated_data/"},
            format="parquet"
        )
    
    print(f"Validation completed. Results: {json.dumps(results, indent=2)}")

except Exception as e:
    print(f"Validation job failed: {e}")
    raise

job.commit()
'''
        
        return validation_script
    
    def implement_data_lineage_tracking(self, 
                                        dataset_id: str,
                                        source_location: str,
                                        transformation_details: Dict) -> Dict:
        """Implement comprehensive data lineage tracking"""
        
        lineage_record = {
            'dataset_id': dataset_id,
            'timestamp': datetime.utcnow().isoformat(),
            'source_location': source_location,
            'transformation_details': transformation_details,
            'data_hash': self._calculate_data_hash(source_location),
            'access_log': [],
            'quality_metrics': {},
            'security_classification': transformation_details.get('classification', 'internal')
        }
        
        # Store lineage record in DynamoDB or S3
        lineage_table = 'ml-data-lineage'
        
        try:
            # In a real implementation, you would store this in DynamoDB
            # For this example, we'll store in S3
            lineage_key = f"lineage/{dataset_id}/{datetime.utcnow().strftime('%Y/%m/%d')}/{dataset_id}.json"
            
            self.s3.put_object(
                Bucket='ml-security-metadata',
                Key=lineage_key,
                Body=json.dumps(lineage_record, indent=2),
                ServerSideEncryption='aws:kms',
                SSEKMSKeyId='arn:aws:kms:us-east-1:123456789012:key/12345678-1234-1234-1234-123456789012'
            )
            
            return {
                'lineage_id': dataset_id,
                'lineage_location': f's3://ml-security-metadata/{lineage_key}',
                'status': 'recorded'
            }
            
        except Exception as e:
            print(f"Error recording data lineage: {e}")
            raise
    
    def _calculate_data_hash(self, data_location: str) -> str:
        """Calculate hash of data for integrity verification"""
        
        try:
            # For S3 objects, use ETag as a quick integrity check
            # In production, implement more robust hash calculation
            bucket, key = data_location.replace('s3://', '').split('/', 1)
            
            response = self.s3.head_object(Bucket=bucket, Key=key)
            etag = response.get('ETag', '').strip('"')
            
            return etag
            
        except Exception as e:
            print(f"Error calculating data hash: {e}")
            return hashlib.sha256(data_location.encode()).hexdigest()
    
    def implement_privacy_protection(self, 
                                     data_location: str,
                                     privacy_config: Dict,
                                     output_location: str) -> Dict:
        """Implement privacy protection for sensitive data"""
        
        privacy_job_name = f"privacy-protection-{datetime.utcnow().strftime('%Y%m%d-%H%M%S')}"
        
        # Create Glue job for privacy protection
        job_definition = {
            'Name': privacy_job_name,
            'Role': 'arn:aws:iam::123456789012:role/GluePrivacyProtectionRole',
            'Command': {
                'Name': 'glueetl',
                'ScriptLocation': 's3://ml-security-scripts/privacy_protection.py',
                'PythonVersion': '3'
            },
            'DefaultArguments': {
                '--job-language': 'python',
                '--data_location': data_location,
                '--privacy_config': json.dumps(privacy_config),
                '--output_location': output_location,
                '--enable-metrics': '',
                '--enable-continuous-cloudwatch-log': 'true'
            },
            'MaxRetries': 1,
            'Timeout': 2880,
            'MaxCapacity': 10.0,
            'SecurityConfiguration': 'ml-security-configuration'
        }
        
        try:
            response = self.glue.create_job(**job_definition)
            
            job_run_response = self.glue.start_job_run(
                JobName=privacy_job_name,
                Arguments={
                    '--protection_timestamp': datetime.utcnow().isoformat()
                }
            )
            
            return {
                'job_name': privacy_job_name,
                'job_run_id': job_run_response['JobRunId'],
                'privacy_techniques': privacy_config.get('techniques', []),
                'status': 'started'
            }
            
        except Exception as e:
            print(f"Error creating privacy protection job: {e}")
            raise

Secure Model Training and Development

Training Environment Security

Securing the model training environment is crucial for protecting intellectual property and preventing data exfiltration.

Secure SageMaker Training Implementation

import boto3
import json
from datetime import datetime, timedelta
from typing import Dict, List, Optional

class SecureModelTraining:
    """
    Secure model training implementation for AWS SageMaker
    Provides comprehensive security controls for ML training workflows
    """
    
    def __init__(self, region_name: str = 'us-east-1'):
        self.sagemaker = boto3.client('sagemaker', region_name=region_name)
        self.s3 = boto3.client('s3', region_name=region_name)
        self.kms = boto3.client('kms', region_name=region_name)
        self.iam = boto3.client('iam', region_name=region_name)
        self.ec2 = boto3.client('ec2', region_name=region_name)
        
    def create_secure_training_environment(self, 
                                           training_config: Dict) -> Dict:
        """Create isolated and secure training environment"""
        
        environment_name = training_config.get('environment_name', 
                                               f"secure-training-{datetime.utcnow().strftime('%Y%m%d')}")
        
        # Create VPC for training isolation
        vpc_config = self._create_training_vpc(environment_name)
        
        # Create security groups
        security_groups = self._create_training_security_groups(
            vpc_config['vpc_id'], 
            training_config.get('allowed_ports', [])
        )
        
        # Create IAM roles with minimal permissions
        training_role = self._create_training_role(
            environment_name,
            training_config.get('data_sources', []),
            training_config.get('output_location', '')
        )
        
        # Create KMS key for training encryption
        kms_key = self._create_training_kms_key(environment_name)
        
        return {
            'environment_name': environment_name,
            'vpc_config': vpc_config,
            'security_groups': security_groups,
            'training_role_arn': training_role,
            'kms_key_id': kms_key,
            'created_at': datetime.utcnow().isoformat()
        }
    
    def _create_training_vpc(self, environment_name: str) -> Dict:
        """Create isolated VPC for training environment"""
        
        # Create VPC
        vpc_response = self.ec2.create_vpc(
            CidrBlock='10.0.0.0/16',
            TagSpecifications=[
                {
                    'ResourceType': 'vpc',
                    'Tags': [
                        {'Key': 'Name', 'Value': f'{environment_name}-vpc'},
                        {'Key': 'Environment', 'Value': 'ml-training'},
                        {'Key': 'SecurityLevel', 'Value': 'high'}
                    ]
                }
            ]
        )
        
        vpc_id = vpc_response['Vpc']['VpcId']
        
        # Create private subnets
        private_subnets = []
        availability_zones = self.ec2.describe_availability_zones()['AvailabilityZones']
        
        for i, az in enumerate(availability_zones[:2]):  # Use first 2 AZs
            subnet_response = self.ec2.create_subnet(
                VpcId=vpc_id,
                CidrBlock=f'10.0.{i+1}.0/24',
                AvailabilityZone=az['ZoneName'],
                TagSpecifications=[
                    {
                        'ResourceType': 'subnet',
                        'Tags': [
                            {'Key': 'Name', 'Value': f'{environment_name}-private-subnet-{i+1}'},
                            {'Key': 'Type', 'Value': 'private'}
                        ]
                    }
                ]
            )
            private_subnets.append(subnet_response['Subnet']['SubnetId'])
        
        # Create NAT Gateway for outbound internet access
        # First create public subnet for NAT Gateway
        public_subnet_response = self.ec2.create_subnet(
            VpcId=vpc_id,
            CidrBlock='10.0.100.0/24',
            AvailabilityZone=availability_zones[0]['ZoneName'],
            TagSpecifications=[
                {
                    'ResourceType': 'subnet',
                    'Tags': [
                        {'Key': 'Name', 'Value': f'{environment_name}-public-subnet'},
                        {'Key': 'Type', 'Value': 'public'}
                    ]
                }
            ]
        )
        
        public_subnet_id = public_subnet_response['Subnet']['SubnetId']
        
        # Create Internet Gateway
        igw_response = self.ec2.create_internet_gateway(
            TagSpecifications=[
                {
                    'ResourceType': 'internet-gateway',
                    'Tags': [
                        {'Key': 'Name', 'Value': f'{environment_name}-igw'}
                    ]
                }
            ]
        )
        
        igw_id = igw_response['InternetGateway']['InternetGatewayId']
        
        # Attach IGW to VPC
        self.ec2.attach_internet_gateway(
            InternetGatewayId=igw_id,
            VpcId=vpc_id
        )
        
        # Allocate Elastic IP for NAT Gateway
        eip_response = self.ec2.allocate_address(Domain='vpc')
        
        # Create NAT Gateway
        nat_response = self.ec2.create_nat_gateway(
            SubnetId=public_subnet_id,
            AllocationId=eip_response['AllocationId'],
            TagSpecifications=[
                {
                    'ResourceType': 'nat-gateway',
                    'Tags': [
                        {'Key': 'Name', 'Value': f'{environment_name}-nat'}
                    ]
                }
            ]
        )
        
        # Create route tables
        # Public route table
        public_rt_response = self.ec2.create_route_table(
            VpcId=vpc_id,
            TagSpecifications=[
                {
                    'ResourceType': 'route-table',
                    'Tags': [
                        {'Key': 'Name', 'Value': f'{environment_name}-public-rt'}
                    ]
                }
            ]
        )
        
        # Add route to internet gateway
        self.ec2.create_route(
            RouteTableId=public_rt_response['RouteTable']['RouteTableId'],
            DestinationCidrBlock='0.0.0.0/0',
            GatewayId=igw_id
        )
        
        # Associate public subnet with public route table
        self.ec2.associate_route_table(
            SubnetId=public_subnet_id,
            RouteTableId=public_rt_response['RouteTable']['RouteTableId']
        )
        
        # Private route table
        private_rt_response = self.ec2.create_route_table(
            VpcId=vpc_id,
            TagSpecifications=[
                {
                    'ResourceType': 'route-table',
                    'Tags': [
                        {'Key': 'Name', 'Value': f'{environment_name}-private-rt'}
                    ]
                }
            ]
        )
        
        # Add route to NAT gateway (wait for NAT to be available)
        # In production, you should wait for NAT gateway to be available
        self.ec2.create_route(
            RouteTableId=private_rt_response['RouteTable']['RouteTableId'],
            DestinationCidrBlock='0.0.0.0/0',
            NatGatewayId=nat_response['NatGateway']['NatGatewayId']
        )
        
        # Associate private subnets with private route table
        for subnet_id in private_subnets:
            self.ec2.associate_route_table(
                SubnetId=subnet_id,
                RouteTableId=private_rt_response['RouteTable']['RouteTableId']
            )
        
        return {
            'vpc_id': vpc_id,
            'private_subnets': private_subnets,
            'public_subnet': public_subnet_id,
            'nat_gateway_id': nat_response['NatGateway']['NatGatewayId'],
            'internet_gateway_id': igw_id
        }
    
    def _create_training_security_groups(self, vpc_id: str, allowed_ports: List[int]) -> Dict:
        """Create security groups for training environment"""
        
        # Training instances security group
        training_sg_response = self.ec2.create_security_group(
            GroupName='ml-training-instances',
            Description='Security group for ML training instances',
            VpcId=vpc_id,
            TagSpecifications=[
                {
                    'ResourceType': 'security-group',
                    'Tags': [
                        {'Key': 'Name', 'Value': 'ml-training-instances'},
                        {'Key': 'Purpose', 'Value': 'ML Training'}
                    ]
                }
            ]
        )
        
        training_sg_id = training_sg_response['GroupId']
        
        # Allow HTTPS outbound
        self.ec2.authorize_security_group_egress(
            GroupId=training_sg_id,
            IpPermissions=[
                {
                    'IpProtocol': 'tcp',
                    'FromPort': 443,
                    'ToPort': 443,
                    'IpRanges': [{'CidrIp': '0.0.0.0/0'}]
                }
            ]
        )
        
        # Allow custom ports if specified
        for port in allowed_ports:
            self.ec2.authorize_security_group_ingress(
                GroupId=training_sg_id,
                IpPermissions=[
                    {
                        'IpProtocol': 'tcp',
                        'FromPort': port,
                        'ToPort': port,
                        'UserIdGroupPairs': [{'GroupId': training_sg_id}]
                    }
                ]
            )
        
        # VPC Endpoints security group
        vpc_endpoint_sg_response = self.ec2.create_security_group(
            GroupName='ml-vpc-endpoints',
            Description='Security group for VPC endpoints',
            VpcId=vpc_id,
            TagSpecifications=[
                {
                    'ResourceType': 'security-group',
                    'Tags': [
                        {'Key': 'Name', 'Value': 'ml-vpc-endpoints'},
                        {'Key': 'Purpose', 'Value': 'VPC Endpoints'}
                    ]
                }
            ]
        )
        
        vpc_endpoint_sg_id = vpc_endpoint_sg_response['GroupId']
        
        # Allow HTTPS from training instances
        self.ec2.authorize_security_group_ingress(
            GroupId=vpc_endpoint_sg_id,
            IpPermissions=[
                {
                    'IpProtocol': 'tcp',
                    'FromPort': 443,
                    'ToPort': 443,
                    'UserIdGroupPairs': [{'GroupId': training_sg_id}]
                }
            ]
        )
        
        return {
            'training_security_group_id': training_sg_id,
            'vpc_endpoint_security_group_id': vpc_endpoint_sg_id
        }
    
    def _create_training_role(self, 
                              environment_name: str,
                              data_sources: List[str],
                              output_location: str) -> str:
        """Create IAM role with minimal permissions for training"""
        
        role_name = f'{environment_name}-training-role'
        
        # Trust policy for SageMaker
        trust_policy = {
            "Version": "2012-10-17",
            "Statement": [
                {
                    "Effect": "Allow",
                    "Principal": {
                        "Service": "sagemaker.amazonaws.com"
                    },
                    "Action": "sts:AssumeRole",
                    "Condition": {
                        "StringEquals": {
                            "aws:RequestedRegion": boto3.Session().region_name
                        }
                    }
                }
            ]
        }
        
        # Create role
        role_response = self.iam.create_role(
            RoleName=role_name,
            AssumeRolePolicyDocument=json.dumps(trust_policy),
            Description=f'Training role for {environment_name}',
            MaxSessionDuration=3600,  # 1 hour maximum
            Tags=[
                {'Key': 'Environment', 'Value': environment_name},
                {'Key': 'Purpose', 'Value': 'ML Training'}
            ]
        )
        
        # Create minimal permission policy
        s3_resources = []
        for source in data_sources:
            if source.startswith('s3://'):
                bucket_and_prefix = source.replace('s3://', '')
                s3_resources.extend([
                    f'arn:aws:s3:::{bucket_and_prefix}',
                    f'arn:aws:s3:::{bucket_and_prefix}/*'
                ])
        
        if output_location.startswith('s3://'):
            bucket_and_prefix = output_location.replace('s3://', '')
            s3_resources.extend([
                f'arn:aws:s3:::{bucket_and_prefix}',
                f'arn:aws:s3:::{bucket_and_prefix}/*'
            ])
        
        permission_policy = {
            "Version": "2012-10-17",
            "Statement": [
                {
                    "Effect": "Allow",
                    "Action": [
                        "s3:GetObject",
                        "s3:ListBucket"
                    ],
                    "Resource": s3_resources,
                    "Condition": {
                        "StringEquals": {
                            "s3:x-amz-server-side-encryption": "aws:kms"
                        }
                    }
                },
                {
                    "Effect": "Allow",
                    "Action": [
                        "s3:PutObject"
                    ],
                    "Resource": [r for r in s3_resources if r.endswith('/*')],
                    "Condition": {
                        "StringEquals": {
                            "s3:x-amz-server-side-encryption": "aws:kms"
                        }
                    }
                },
                {
                    "Effect": "Allow",
                    "Action": [
                        "kms:Decrypt",
                        "kms:DescribeKey",
                        "kms:Encrypt",
                        "kms:GenerateDataKey",
                        "kms:ReEncrypt*"
                    ],
                    "Resource": "*",
                    "Condition": {
                        "StringEquals": {
                            "kms:ViaService": f"s3.{boto3.Session().region_name}.amazonaws.com"
                        }
                    }
                },
                {
                    "Effect": "Allow",
                    "Action": [
                        "logs:CreateLogGroup",
                        "logs:CreateLogStream",
                        "logs:PutLogEvents"
                    ],
                    "Resource": f"arn:aws:logs:{boto3.Session().region_name}:*:log-group:/aws/sagemaker/*"
                },
                {
                    "Effect": "Allow",
                    "Action": [
                        "cloudwatch:PutMetricData"
                    ],
                    "Resource": "*",
                    "Condition": {
                        "StringEquals": {
                            "cloudwatch:namespace": "AWS/SageMaker"
                        }
                    }
                }
            ]
        }
        
        policy_name = f'{environment_name}-training-policy'
        
        self.iam.put_role_policy(
            RoleName=role_name,
            PolicyName=policy_name,
            PolicyDocument=json.dumps(permission_policy)
        )
        
        return role_response['Role']['Arn']
    
    def _create_training_kms_key(self, environment_name: str) -> str:
        """Create KMS key for training encryption"""
        
        key_policy = {
            "Version": "2012-10-17",
            "Statement": [
                {
                    "Sid": "Enable IAM User Permissions",
                    "Effect": "Allow",
                    "Principal": {
                        "AWS": f"arn:aws:iam::{boto3.client('sts').get_caller_identity()['Account']}:root"
                    },
                    "Action": "kms:*",
                    "Resource": "*"
                },
                {
                    "Sid": "Allow SageMaker Service",
                    "Effect": "Allow",
                    "Principal": {
                        "Service": "sagemaker.amazonaws.com"
                    },
                    "Action": [
                        "kms:Decrypt",
                        "kms:DescribeKey",
                        "kms:Encrypt",
                        "kms:GenerateDataKey*",
                        "kms:ReEncrypt*"
                    ],
                    "Resource": "*"
                }
            ]
        }
        
        key_response = self.kms.create_key(
            Policy=json.dumps(key_policy),
            Description=f'KMS key for {environment_name} ML training',
            Usage='ENCRYPT_DECRYPT',
            KeySpec='SYMMETRIC_DEFAULT',
            Tags=[
                {'TagKey': 'Environment', 'TagValue': environment_name},
                {'TagKey': 'Purpose', 'TagValue': 'ML Training Encryption'}
            ]
        )
        
        # Create alias
        self.kms.create_alias(
            AliasName=f'alias/{environment_name}-training-key',
            TargetKeyId=key_response['KeyMetadata']['KeyId']
        )
        
        return key_response['KeyMetadata']['KeyId']
    
    def launch_secure_training_job(self, 
                                   job_config: Dict,
                                   environment_config: Dict) -> Dict:
        """Launch secure training job with comprehensive protection"""
        
        job_name = job_config.get('job_name', 
                                  f"secure-training-{datetime.utcnow().strftime('%Y%m%d-%H%M%S')}")
        
        # Prepare VPC configuration
        vpc_config = {
            'SecurityGroupIds': [environment_config['security_groups']['training_security_group_id']],
            'Subnets': environment_config['vpc_config']['private_subnets']
        }
        
        # Training job configuration with security controls
        training_job_config = {
            'TrainingJobName': job_name,
            'RoleArn': environment_config['training_role_arn'],
            'AlgorithmSpecification': {
                'TrainingImage': job_config['training_image'],
                'TrainingInputMode': 'File',
                'EnableSageMakerMetricsTimeSeries': True
            },
            'InputDataConfig': self._prepare_secure_input_config(
                job_config['input_data'],
                environment_config['kms_key_id']
            ),
            'OutputDataConfig': {
                'S3OutputPath': job_config['output_location'],
                'KmsKeyId': environment_config['kms_key_id']
            },
            'ResourceConfig': {
                'InstanceType': job_config.get('instance_type', 'ml.m5.large'),
                'InstanceCount': job_config.get('instance_count', 1),
                'VolumeSizeInGB': job_config.get('volume_size', 30),
                'VolumeKmsKeyId': environment_config['kms_key_id']
            },
            'VpcConfig': vpc_config,
            'StoppingCondition': {
                'MaxRuntimeInSeconds': job_config.get('max_runtime', 86400)  # 24 hours default
            },
            'EnableNetworkIsolation': True,
            'EnableInterContainerTrafficEncryption': True,
            'EnableManagedSpotTraining': False,  # Disable for security
            'HyperParameters': self._sanitize_hyperparameters(
                job_config.get('hyperparameters', {})
            ),
            'Tags': [
                {'Key': 'Environment', 'Value': environment_config['environment_name']},
                {'Key': 'SecurityLevel', 'Value': 'high'},
                {'Key': 'DataClassification', 'Value': job_config.get('data_classification', 'confidential')},
                {'Key': 'Owner', 'Value': job_config.get('owner', 'ml-team')}
            ],
            'ExperimentConfig': {
                'ExperimentName': f"{job_config.get('experiment_name', job_name)}-experiment"
            },
            'DebugHookConfig': {
                'S3OutputPath': f"{job_config['output_location']}/debug",
                'HookParameters': {
                    'save_interval': '100',
                    'include_regex': '.*gradient.*|.*weight.*|.*bias.*'
                }
            },
            'ProfilerConfig': {
                'S3OutputPath': f"{job_config['output_location']}/profiler",
                'ProfilingIntervalInMilliseconds': 500,
                'ProfilingParameters': {
                    'DataloaderProfilingConfig': '{"StartStep": 5, "NumSteps": 3}',
                    'DetailedProfilingConfig': '{"StartStep": 5, "NumSteps": 3}',
                    'PythonProfilingConfig': '{"StartStep": 5, "NumSteps": 3}'
                }
            },
            'TensorBoardOutputConfig': {
                'S3OutputPath': f"{job_config['output_location']}/tensorboard",
                'LocalPath': '/opt/ml/output/tensorboard'
            }
        }
        
        try:
            response = self.sagemaker.create_training_job(**training_job_config)
            
            # Set up training job monitoring
            monitoring_config = self._setup_training_monitoring(job_name, environment_config)
            
            return {
                'training_job_name': job_name,
                'training_job_arn': response['TrainingJobArn'],
                'environment_name': environment_config['environment_name'],
                'security_features': [
                    'network_isolation_enabled',
                    'encryption_in_transit',
                    'encryption_at_rest',
                    'vpc_isolation',
                    'minimal_iam_permissions',
                    'security_monitoring_enabled'
                ],
                'monitoring_config': monitoring_config,
                'status': 'started'
            }
            
        except Exception as e:
            print(f"Error launching secure training job: {e}")
            raise
    
    def _prepare_secure_input_config(self, input_data: List[Dict], kms_key_id: str) -> List[Dict]:
        """Prepare secure input data configuration"""
        
        input_config = []
        
        for data_source in input_data:
            config = {
                'ChannelName': data_source['channel_name'],
                'DataSource': {
                    'S3DataSource': {
                        'S3DataType': data_source.get('data_type', 'S3Prefix'),
                        'S3Uri': data_source['s3_uri'],
                        'S3DataDistributionType': data_source.get('distribution_type', 'FullyReplicated')
                    }
                },
                'ContentType': data_source.get('content_type', 'application/json'),
                'CompressionType': data_source.get('compression', 'None'),
                'InputMode': data_source.get('input_mode', 'File')
            }
            
            # Add encryption configuration
            if kms_key_id:
                config['DataSource']['S3DataSource']['KmsKeyId'] = kms_key_id
            
            input_config.append(config)
        
        return input_config
    
    def _sanitize_hyperparameters(self, hyperparameters: Dict) -> Dict:
        """Sanitize hyperparameters to prevent information leakage"""
        
        sanitized = {}
        
        # List of sensitive parameter names to exclude
        sensitive_params = [
            'password', 'secret', 'token', 'key', 'credential',
            'api_key', 'access_key', 'private_key'
        ]
        
        for key, value in hyperparameters.items():
            # Check if parameter name is sensitive
            if any(sensitive in key.lower() for sensitive in sensitive_params):
                continue
            
            # Convert all values to strings (SageMaker requirement)
            sanitized[key] = str(value)
        
        return sanitized
    
    def _setup_training_monitoring(self, job_name: str, environment_config: Dict) -> Dict:
        """Set up comprehensive monitoring for training job"""
        
        monitoring_config = {
            'cloudwatch_alarms': [],
            'custom_metrics': [],
            'log_monitoring': True
        }
        
        # CloudWatch alarms would be set up here
        # Custom metrics collection would be configured
        # Log monitoring and alerting would be enabled
        
        return monitoring_config

Model Deployment Security

Secure Inference Endpoints

Securing model deployment involves protecting inference endpoints, implementing authentication, and monitoring for adversarial attacks.

Secure Endpoint Deployment

# CloudFormation template for secure SageMaker endpoint deployment
AWSTemplateFormatVersion: '2010-09-09'
Description: 'Secure SageMaker model endpoint with comprehensive security controls'

Parameters:
  ModelName:
    Type: String
    Description: Name of the SageMaker model
  
  EndpointName:
    Type: String
    Description: Name for the inference endpoint
  
  VPCId:
    Type: AWS::EC2::VPC::Id
    Description: VPC ID for endpoint deployment
  
  PrivateSubnetIds:
    Type: List<AWS::EC2::Subnet::Id>
    Description: Private subnet IDs for endpoint
  
  KMSKeyId:
    Type: String
    Description: KMS key ID for encryption
  
  DataClassification:
    Type: String
    Default: confidential
    AllowedValues: [public, internal, confidential, restricted]

Resources:
  # Security group for endpoint
  EndpointSecurityGroup:
    Type: AWS::EC2::SecurityGroup
    Properties:
      GroupDescription: Security group for SageMaker endpoint
      VpcId: !Ref VPCId
      SecurityGroupIngress:
        - IpProtocol: tcp
          FromPort: 443
          ToPort: 443
          SourceSecurityGroupId: !Ref ClientSecurityGroup
      SecurityGroupEgress:
        - IpProtocol: tcp
          FromPort: 443
          ToPort: 443
          CidrIp: 0.0.0.0/0
      Tags:
        - Key: Name
          Value: !Sub '${EndpointName}-endpoint-sg'
        - Key: Purpose
          Value: ML Inference Endpoint

  # Security group for clients
  ClientSecurityGroup:
    Type: AWS::EC2::SecurityGroup
    Properties:
      GroupDescription: Security group for endpoint clients
      VpcId: !Ref VPCId
      SecurityGroupEgress:
        - IpProtocol: tcp
          FromPort: 443
          ToPort: 443
          DestinationSecurityGroupId: !Ref EndpointSecurityGroup
      Tags:
        - Key: Name
          Value: !Sub '${EndpointName}-client-sg'

  # IAM role for endpoint
  EndpointExecutionRole:
    Type: AWS::IAM::Role
    Properties:
      RoleName: !Sub '${EndpointName}-execution-role'
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal:
              Service: sagemaker.amazonaws.com
            Action: sts:AssumeRole
            Condition:
              StringEquals:
                'aws:RequestedRegion': !Ref AWS::Region
      ManagedPolicyArns:
        - arn:aws:iam::aws:policy/AmazonSageMakerFullAccess
      Policies:
        - PolicyName: EndpointKMSAccess
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: Allow
                Action:
                  - kms:Decrypt
                  - kms:DescribeKey
                Resource: !Sub 'arn:aws:kms:${AWS::Region}:${AWS::AccountId}:key/${KMSKeyId}'
        - PolicyName: EndpointLogging
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: Allow
                Action:
                  - logs:CreateLogGroup
                  - logs:CreateLogStream
                  - logs:PutLogEvents
                Resource: !Sub 'arn:aws:logs:${AWS::Region}:${AWS::AccountId}:*'

  # VPC endpoint for SageMaker runtime
  SageMakerRuntimeVPCEndpoint:
    Type: AWS::EC2::VPCEndpoint
    Properties:
      VpcId: !Ref VPCId
      ServiceName: !Sub 'com.amazonaws.${AWS::Region}.sagemaker.runtime'
      VpcEndpointType: Interface
      SubnetIds: !Ref PrivateSubnetIds
      SecurityGroupIds:
        - !Ref EndpointSecurityGroup
      PolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal: '*'
            Action:
              - sagemaker:InvokeEndpoint
            Resource: '*'
            Condition:
              StringEquals:
                'aws:PrincipalTag/Environment': production

  # Endpoint configuration
  EndpointConfig:
    Type: AWS::SageMaker::EndpointConfig
    Properties:
      EndpointConfigName: !Sub '${EndpointName}-config'
      ProductionVariants:
        - VariantName: primary
          ModelName: !Ref ModelName
          InitialInstanceCount: 2
          InstanceType: ml.m5.large
          InitialVariantWeight: 1.0
          AcceleratorType: !Ref AWS::NoValue
      DataCaptureConfig:
        EnableCapture: true
        InitialSamplingPercentage: 20
        DestinationS3Uri: !Sub 's3://ml-inference-capture-${AWS::AccountId}/endpoints/${EndpointName}/'
        KmsKeyId: !Ref KMSKeyId
        CaptureOptions:
          - CaptureMode: Input
          - CaptureMode: Output
        CaptureContentTypeHeader:
          CsvContentTypes:
            - text/csv
          JsonContentTypes:
            - application/json
      Tags:
        - Key: Environment
          Value: production
        - Key: DataClassification
          Value: !Ref DataClassification
        - Key: SecurityLevel
          Value: high
      KmsKeyId: !Ref KMSKeyId
      AsyncInferenceConfig:
        OutputConfig:
          S3OutputPath: !Sub 's3://ml-async-inference-${AWS::AccountId}/${EndpointName}/output/'
          KmsKeyId: !Ref KMSKeyId
          NotificationConfig:
            SuccessTopic: !Ref InferenceSuccessTopic
            ErrorTopic: !Ref InferenceErrorTopic
        ClientConfig:
          MaxConcurrentInvocationsPerInstance: 4

  # SageMaker endpoint
  Endpoint:
    Type: AWS::SageMaker::Endpoint
    Properties:
      EndpointName: !Ref EndpointName
      EndpointConfigName: !Ref EndpointConfig
      Tags:
        - Key: Name
          Value: !Ref EndpointName
        - Key: Environment
          Value: production
        - Key: DataClassification
          Value: !Ref DataClassification

  # SNS topics for notifications
  InferenceSuccessTopic:
    Type: AWS::SNS::Topic
    Properties:
      TopicName: !Sub '${EndpointName}-inference-success'
      DisplayName: Inference Success Notifications
      KmsMasterKeyId: !Ref KMSKeyId

  InferenceErrorTopic:
    Type: AWS::SNS::Topic
    Properties:
      TopicName: !Sub '${EndpointName}-inference-error'
      DisplayName: Inference Error Notifications
      KmsMasterKeyId: !Ref KMSKeyId

  # CloudWatch log group for endpoint
  EndpointLogGroup:
    Type: AWS::Logs::LogGroup
    Properties:
      LogGroupName: !Sub '/aws/sagemaker/Endpoints/${EndpointName}'
      RetentionInDays: 90
      KmsKeyId: !Sub 'arn:aws:kms:${AWS::Region}:${AWS::AccountId}:key/${KMSKeyId}'

  # CloudWatch alarms for monitoring
  EndpointInvocationsAlarm:
    Type: AWS::CloudWatch::Alarm
    Properties:
      AlarmName: !Sub '${EndpointName}-high-invocations'
      AlarmDescription: High number of endpoint invocations
      MetricName: Invocations
      Namespace: AWS/SageMaker
      Statistic: Sum
      Period: 300
      EvaluationPeriods: 2
      Threshold: 1000
      ComparisonOperator: GreaterThanThreshold
      Dimensions:
        - Name: EndpointName
          Value: !Ref EndpointName
      AlarmActions:
        - !Ref InferenceErrorTopic

  EndpointLatencyAlarm:
    Type: AWS::CloudWatch::Alarm
    Properties:
      AlarmName: !Sub '${EndpointName}-high-latency'
      AlarmDescription: High endpoint latency
      MetricName: ModelLatency
      Namespace: AWS/SageMaker
      Statistic: Average
      Period: 300
      EvaluationPeriods: 2
      Threshold: 10000  # 10 seconds
      ComparisonOperator: GreaterThanThreshold
      Dimensions:
        - Name: EndpointName
          Value: !Ref EndpointName
      AlarmActions:
        - !Ref InferenceErrorTopic

  EndpointErrorRateAlarm:
    Type: AWS::CloudWatch::Alarm
    Properties:
      AlarmName: !Sub '${EndpointName}-high-error-rate'
      AlarmDescription: High endpoint error rate
      MetricName: Invocation4XXErrors
      Namespace: AWS/SageMaker
      Statistic: Sum
      Period: 300
      EvaluationPeriods: 1
      Threshold: 10
      ComparisonOperator: GreaterThanThreshold
      Dimensions:
        - Name: EndpointName
          Value: !Ref EndpointName
      AlarmActions:
        - !Ref InferenceErrorTopic

  # WAF for endpoint protection
  EndpointWebACL:
    Type: AWS::WAFv2::WebACL
    Properties:
      Name: !Sub '${EndpointName}-waf'
      Scope: REGIONAL
      DefaultAction:
        Allow: {}
      Rules:
        - Name: RateLimitRule
          Priority: 1
          Statement:
            RateBasedStatement:
              Limit: 2000
              AggregateKeyType: IP
          Action:
            Block: {}
          VisibilityConfig:
            SampledRequestsEnabled: true
            CloudWatchMetricsEnabled: true
            MetricName: !Sub '${EndpointName}-rate-limit'
        - Name: IPReputationRule
          Priority: 2
          Statement:
            ManagedRuleGroupStatement:
              VendorName: AWS
              Name: AWSManagedRulesAmazonIpReputationList
          OverrideAction:
            None: {}
          VisibilityConfig:
            SampledRequestsEnabled: true
            CloudWatchMetricsEnabled: true
            MetricName: !Sub '${EndpointName}-ip-reputation'
        - Name: KnownBadInputsRule
          Priority: 3
          Statement:
            ManagedRuleGroupStatement:
              VendorName: AWS
              Name: AWSManagedRulesKnownBadInputsRuleSet
          OverrideAction:
            None: {}
          VisibilityConfig:
            SampledRequestsEnabled: true
            CloudWatchMetricsEnabled: true
            MetricName: !Sub '${EndpointName}-bad-inputs'

Outputs:
  EndpointName:
    Description: Name of the created endpoint
    Value: !Ref Endpoint
    Export:
      Name: !Sub '${AWS::StackName}-endpoint-name'
  
  EndpointArn:
    Description: ARN of the created endpoint
    Value: !Ref Endpoint
    Export:
      Name: !Sub '${AWS::StackName}-endpoint-arn'
  
  EndpointUrl:
    Description: URL for invoking the endpoint
    Value: !Sub 'https://runtime.sagemaker.${AWS::Region}.amazonaws.com/endpoints/${EndpointName}/invocations'
    Export:
      Name: !Sub '${AWS::StackName}-endpoint-url'
  
  SecurityGroupId:
    Description: Security group ID for endpoint clients
    Value: !Ref ClientSecurityGroup
    Export:
      Name: !Sub '${AWS::StackName}-client-sg-id'

Implementation Roadmap for Secure ML Pipelines

Phase 1: Data Security Foundation (Weeks 1-3)

Week 1: Data Infrastructure Security

• Implement secure S3 buckets with KMS encryption
• Configure VPC endpoints for data access
• Set up IAM policies with least privilege access
• Deploy data classification and tagging system

Week 2: Data Pipeline Security

• Implement data validation and quality checks
• Deploy PII detection and anonymization
• Set up data lineage tracking system
• Configure audit logging for data access

Week 3: Privacy Protection Implementation

• Deploy differential privacy techniques
• Implement data masking and anonymization
• Set up consent management system
• Configure privacy impact assessments

Phase 2: Training Security (Weeks 4-6)

Week 4: Training Environment Setup

• Create isolated VPCs for training
• Deploy network security controls
• Implement container security scanning
• Set up secure training image repositories

Week 5: Model Security Implementation

• Deploy model encryption and signing
• Implement model versioning and integrity checks
• Set up model artifact protection
• Configure training job monitoring

Week 6: Advanced Training Security

• Implement federated learning capabilities
• Deploy secure multi-party computation
• Set up homomorphic encryption for training
• Configure adversarial training techniques

Phase 3: Deployment Security (Weeks 7-9)

Week 7: Inference Security

• Deploy secure model endpoints with authentication
• Implement rate limiting and WAF protection
• Set up inference monitoring and alerting
• Configure model performance tracking

Week 8: Production Hardening

• Implement A/B testing security controls
• Deploy canary deployment automation
• Set up model rollback capabilities
• Configure production incident response

Week 9: Advanced Deployment Security

• Implement adversarial attack detection
• Deploy model explanation and interpretability
• Set up bias detection and mitigation
• Configure continuous security validation

Phase 4: Monitoring and Governance (Weeks 10-12)

Week 10: Security Monitoring

• Deploy SIEM integration for ML events
• Implement behavioral analytics for users
• Set up automated threat detection
• Configure security dashboards and alerting

Week 11: Compliance and Governance

• Implement compliance monitoring automation
• Deploy policy-as-code for ML governance
• Set up audit reporting and evidence collection
• Configure risk assessment automation

Week 12: Optimization and Maturity

• Conduct security maturity assessment
• Optimize performance and cost efficiency
• Implement advanced threat hunting
• Establish continuous improvement processes

Related Articles and Additional Resources

AWS Documentation

• SageMaker Security Best Practices
• MLOps Security Guide
• Data Pipeline Security

Industry Standards

• NIST AI Risk Management Framework
• ISO/IEC 23053:2022 AI Security
• OWASP ML Security Top 10

Community Resources

• MLSecOps Community
• AWS ML Security Samples
• SageMaker Security Workshop

---

This comprehensive guide provides the foundation for implementing secure AI/ML pipelines on AWS. The combination of data protection, model security, and deployment safety creates a robust MLOps security posture for enterprise environments.