MLSecOps Guide Secure ML Pipelines Production 2026

# ml_security/data_validator.py
import pandas as pd
import numpy as np
from typing import Dict, List, Optional
from scipy import stats
from sklearn.ensemble import IsolationForest
import logging
from dataclasses import dataclass
from datetime import datetime

logger = logging.getLogger(__name__)

@dataclass
class ValidationResult:
    """Result of data validation check"""
    passed: bool
    anomalies_detected: int
    drift_score: float
    flagged_samples: List[int]
    timestamp: datetime

class MLDataValidator:
    """
    Comprehensive data validator for ML pipelines
    Detects poisoning, drift, and anomalies
    """

    def __init__(
        self,
        reference_data: pd.DataFrame,
        contamination_threshold: float = 0.05,
        drift_threshold: float = 0.15
    ):
        self.reference_data = reference_data
        self.contamination_threshold = contamination_threshold
        self.drift_threshold = drift_threshold

        # Fit isolation forest on clean reference data
        self.isolation_forest = IsolationForest(
            contamination=contamination_threshold,
            random_state=42
        )
        self.isolation_forest.fit(reference_data.select_dtypes(include=[np.number]))

        # Compute reference statistics
        self.reference_stats = self._compute_statistics(reference_data)

        logger.info(f"Data validator initialized with {len(reference_data)} reference samples")

    def validate(self, new_data: pd.DataFrame) -> ValidationResult:
        """
        Validate new data batch for poisoning and drift
        """
        start_time = datetime.now()

        # 1. Schema validation
        if not self._validate_schema(new_data):
            logger.error("Schema validation failed")
            return ValidationResult(
                passed=False,
                anomalies_detected=len(new_data),
                drift_score=1.0,
                flagged_samples=list(range(len(new_data))),
                timestamp=start_time
            )

        # 2. Anomaly detection (potential poisoning)
        numeric_data = new_data.select_dtypes(include=[np.number])
        anomaly_predictions = self.isolation_forest.predict(numeric_data)
        anomaly_indices = np.where(anomaly_predictions == -1)[0].tolist()

        # 3. Statistical drift detection
        drift_score = self._compute_drift(new_data)

        # 4. Bias detection (check for demographic shifts)
        bias_flagged = self._detect_bias_shift(new_data)

        # Combine results
        flagged_samples = list(set(anomaly_indices + bias_flagged))
        passed = (
            len(flagged_samples) / len(new_data) < self.contamination_threshold
            and drift_score < self.drift_threshold
        )

        result = ValidationResult(
            passed=passed,
            anomalies_detected=len(flagged_samples),
            drift_score=drift_score,
            flagged_samples=flagged_samples,
            timestamp=start_time
        )

        # Log results
        logger.info(
            f"Validation completed: passed={passed}, "
            f"anomalies={len(flagged_samples)}/{len(new_data)}, "
            f"drift_score={drift_score:.3f}"
        )

        return result

    def _validate_schema(self, data: pd.DataFrame) -> bool:
        """Validate data schema matches reference"""
        return (
            set(data.columns) == set(self.reference_data.columns)
            and all(data[col].dtype == self.reference_data[col].dtype
                   for col in data.columns)
        )

    def _compute_statistics(self, data: pd.DataFrame) -> Dict:
        """Compute statistical summary of data"""
        numeric_cols = data.select_dtypes(include=[np.number]).columns
        return {
            col: {
                'mean': data[col].mean(),
                'std': data[col].std(),
                'min': data[col].min(),
                'max': data[col].max(),
                'median': data[col].median()
            }
            for col in numeric_cols
        }

    def _compute_drift(self, new_data: pd.DataFrame) -> float:
        """
        Compute drift using Kolmogorov-Smirnov test
        Returns drift score [0, 1] where >0.15 indicates significant drift
        """
        numeric_cols = new_data.select_dtypes(include=[np.number]).columns
        drift_scores = []

        for col in numeric_cols:
            if col in self.reference_data.columns:
                # KS test compares distributions
                ks_stat, p_value = stats.ks_2samp(
                    self.reference_data[col],
                    new_data[col]
                )
                drift_scores.append(ks_stat)

        return np.mean(drift_scores) if drift_scores else 0.0

    def _detect_bias_shift(self, data: pd.DataFrame) -> List[int]:
        """
        Detect sudden shifts in sensitive attributes (demographic bias)
        """
        flagged = []
        sensitive_cols = ['age', 'gender', 'ethnicity']  # Configure per use case

        for col in sensitive_cols:
            if col in data.columns:
                ref_dist = self.reference_data[col].value_counts(normalize=True)
                new_dist = data[col].value_counts(normalize=True)

                # Chi-square test for distribution shift
                chi2, p_value = stats.chisquare(
                    new_dist.reindex(ref_dist.index, fill_value=0),
                    ref_dist
                )

                if p_value < 0.01:  # Significant shift detected
                    logger.warning(f"Bias shift detected in {col}: p={p_value:.4f}")

        return flagged

# ml_security/model_wrapper.py
import time
import hashlib
from typing import Any, Dict, Optional
from collections import defaultdict
from datetime import datetime, timedelta
import numpy as np
import logging

logger = logging.getLogger(__name__)

class SecureModelWrapper:
    """
    Security wrapper for ML model inference
    Implements rate limiting, adversarial detection, audit logging
    """

    def __init__(
        self,
        model: Any,
        rate_limit: int = 100,  # requests per hour per user
        confidence_precision: int = 2,  # round to 2 decimals
        enable_audit_log: bool = True
    ):
        self.model = model
        self.rate_limit = rate_limit
        self.confidence_precision = confidence_precision
        self.enable_audit_log = enable_audit_log

        # Rate limiting state
        self.request_counts = defaultdict(list)

        # Adversarial detection (simple entropy-based)
        self.input_history = []

        logger.info(f"Secure model wrapper initialized: rate_limit={rate_limit}/hour")

    def predict(
        self,
        input_data: np.ndarray,
        user_id: str,
        request_metadata: Optional[Dict] = None
    ) -> Dict:
        """
        Secure prediction with rate limiting and monitoring
        """
        start_time = time.time()

        # 1. Rate limiting check
        if not self._check_rate_limit(user_id):
            logger.warning(f"Rate limit exceeded for user {user_id}")
            return {
                "error": "Rate limit exceeded",
                "retry_after": self._get_retry_time(user_id)
            }

        # 2. Input sanitization
        sanitized_input = self._sanitize_input(input_data)

        # 3. Adversarial detection
        if self._is_adversarial(sanitized_input):
            logger.warning(f"Potential adversarial input detected from user {user_id}")
            # Don't block, but flag for review
            adversarial_flag = True
        else:
            adversarial_flag = False

        # 4. Model inference
        try:
            raw_prediction = self.model.predict(sanitized_input)

            # 5. Output protection (hide precision to prevent extraction)
            protected_output = self._protect_output(raw_prediction)

            # 6. Audit logging
            if self.enable_audit_log:
                self._log_request(
                    user_id=user_id,
                    input_hash=self._hash_input(input_data),
                    output=protected_output,
                    latency_ms=(time.time() - start_time) * 1000,
                    adversarial_flag=adversarial_flag,
                    metadata=request_metadata
                )

            return {
                "prediction": protected_output,
                "confidence": self._round_confidence(protected_output),
                "request_id": self._generate_request_id(),
                "warning": "Flagged for review" if adversarial_flag else None
            }

        except Exception as e:
            logger.error(f"Model inference failed: {str(e)}")
            return {"error": "Internal error"}

    def _check_rate_limit(self, user_id: str) -> bool:
        """Check if user is within rate limit"""
        now = datetime.now()
        one_hour_ago = now - timedelta(hours=1)

        # Remove old requests
        self.request_counts[user_id] = [
            ts for ts in self.request_counts[user_id]
            if ts > one_hour_ago
        ]

        # Check limit
        if len(self.request_counts[user_id]) >= self.rate_limit:
            return False

        # Record new request
        self.request_counts[user_id].append(now)
        return True

    def _sanitize_input(self, input_data: np.ndarray) -> np.ndarray:
        """
        Sanitize input to defend against adversarial perturbations
        Techniques: clipping, smoothing, quantization
        """
        # Clip outliers
        sanitized = np.clip(input_data, -10, 10)

        # Add small Gaussian noise (smoothing)
        sanitized += np.random.normal(0, 0.01, sanitized.shape)

        # Quantize (reduce precision)
        sanitized = np.round(sanitized, decimals=4)

        return sanitized

    def _is_adversarial(self, input_data: np.ndarray) -> bool:
        """
        Simple adversarial detection using input entropy
        More sophisticated: use adversarial detectors (LID, Mahalanobis)
        """
        # Compute input entropy
        input_flat = input_data.flatten()
        entropy = stats.entropy(np.abs(input_flat) + 1e-10)

        # High entropy may indicate adversarial perturbations
        return entropy > 10.0  # Threshold tuned on validation set

    def _protect_output(self, output: Any) -> Any:
        """Round confidence scores to prevent extraction attacks"""
        if isinstance(output, np.ndarray):
            return np.round(output, decimals=self.confidence_precision)
        return output

    def _round_confidence(self, output: Any) -> float:
        """Extract and round confidence score"""
        if isinstance(output, np.ndarray):
            return round(float(np.max(output)), self.confidence_precision)
        return 0.0

    def _hash_input(self, input_data: np.ndarray) -> str:
        """Hash input for audit trail without storing raw data"""
        return hashlib.sha256(input_data.tobytes()).hexdigest()[:16]

    def _log_request(self, **kwargs):
        """Log request to audit trail (implement with your logging backend)"""
        log_entry = {
            "timestamp": datetime.now().isoformat(),
            **kwargs
        }
        logger.info(f"Audit: {log_entry}")

    def _generate_request_id(self) -> str:
        """Generate unique request ID"""
        return hashlib.sha256(
            f"{datetime.now().isoformat()}{np.random.random()}".encode()
        ).hexdigest()[:12]

    def _get_retry_time(self, user_id: str) -> int:
        """Get seconds until rate limit resets"""
        if not self.request_counts[user_id]:
            return 0
        oldest_request = min(self.request_counts[user_id])
        one_hour_later = oldest_request + timedelta(hours=1)
        return max(0, int((one_hour_later - datetime.now()).total_seconds()))