System Design Interview Questions (DS & ML)

┌─────────────┐
│ User Activity│ (clicks, views, purchases, time spent)
└──────┬──────┘
       ↓
┌─────────────────────────────────────────┐
│        Feature Engineering              │
│  - Real-time: last 1hr behavior         │
│  - Batch: 7d/30d aggregates             │
│  - User profile: demographics, history  │
│  - Context: time, device, location      │
└──────┬──────────────────────────────────┘
       ↓
┌─────────────────────────────────────────┐
│     Candidate Generation (Retrieval)    │
│  ┌────────────────────────────────────┐ │
│  │ 1. Collaborative Filtering (ALS)   │ │ → 1000 candidates
│  │ 2. Content-based (embeddings)      │ │
│  │ 3. Trending/Popular items          │ │
│  │ 4. Graph-based (item2item)         │ │
│  └────────────────────────────────────┘ │
└──────┬──────────────────────────────────┘
       ↓
┌─────────────────────────────────────────┐
│          Ranking (Scoring)              │
│  Two-Tower Neural Network               │
│  - User tower: user embeddings          │
│  - Item tower: item embeddings          │
│  - Features: 100+ features              │
│  - Model: DLRM, DCN, DeepFM             │
└──────┬──────────────────────────────────┘ → Top 100
       ↓
┌─────────────────────────────────────────┐
│         Re-ranking (Filtering)          │
│  - Diversity: avoid similar items       │
│  - Business rules: inventory, policies  │
│  - Explore/exploit: Thompson sampling   │
│  - Deduplication                        │
└──────┬──────────────────────────────────┘ → Top 20
       ↓
┌─────────────┐
│   Results   │
└─────────────┘

class RecommendationSystem:
    def __init__(self):
        self.feature_store = FeatureStore()
        self.candidate_gen = CandidateGenerator()
        self.ranker = TwoTowerRanker()
        self.reranker = Reranker()

    def get_recommendations(self, user_id: str, context: dict) -> List[str]:
        # 1. Feature retrieval (<10ms)
        user_features = self.feature_store.get_user_features(user_id)
        context_features = self._extract_context(context)

        # 2. Candidate generation (<20ms)
        # Retrieve ~1000 candidates from multiple sources
        cf_candidates = self.candidate_gen.collaborative_filter(user_id, k=500)
        content_candidates = self.candidate_gen.content_based(user_features, k=300)
        trending = self.candidate_gen.get_trending(k=200)

        all_candidates = set(cf_candidates + content_candidates + trending)

        # 3. Ranking (<15ms)
        # Score all candidates with neural network
        candidate_features = self.feature_store.get_item_features(all_candidates)
        scores = self.ranker.predict(user_features, candidate_features, context_features)

        top_100 = sorted(zip(all_candidates, scores), key=lambda x: x[1], reverse=True)[:100]

        # 4. Re-ranking (<5ms)
        # Apply business rules and diversification
        final_recs = self.reranker.rerank(
            candidates=top_100,
            diversity_weight=0.3,
            explore_rate=0.1
        )

        return [item_id for item_id, _ in final_recs[:20]]

# Candidate Generation with ANN
class CandidateGenerator:
    def collaborative_filter(self, user_id: str, k: int) -> List[str]:
        """Use Approximate Nearest Neighbors for fast retrieval"""
        user_embedding = self.get_user_embedding(user_id)  # 128-dim vector

        # HNSW index for fast ANN search
        # Search through 10M items in <5ms
        similar_items = self.ann_index.search(user_embedding, k=k)
        return similar_items

# Two-Tower Ranking Model
class TwoTowerRanker:
    def __init__(self):
        self.user_tower = UserTower(input_dim=200, output_dim=128)
        self.item_tower = ItemTower(input_dim=150, output_dim=128)

    def predict(self, user_feats, item_feats, context_feats):
        user_emb = self.user_tower(user_feats)
        item_emb = self.item_tower(item_feats)

        # Dot product for scoring
        scores = torch.matmul(user_emb, item_emb.T)
        return scores

def handle_cold_start(user_id: str, user_data: dict):
    """Strategies for new users/items"""

    # New User:
    if is_new_user(user_id):
        # 1. Use demographic-based recommendations
        recs = get_popular_for_demographic(user_data['age'], user_data['location'])

        # 2. Quick onboarding survey
        preferences = get_user_preferences(user_id)
        recs += content_based_on_preferences(preferences)

        # 3. Thompson sampling for exploration
        recs += explore_diverse_content(explore_rate=0.5)

    # New Item:
    if is_new_item(item_id):
        # 1. Content-based: use item metadata
        similar_items = find_similar_by_content(item_id)

        # 2. Cold start boost in ranking
        boost_score = 0.1  # Temporary boost

        # 3. Show to exploratory users first
        target_users = get_early_adopter_users()

┌──────────────┐
│ Transaction  │ (amount, merchant, location, device, etc.)
└──────┬───────┘
       ↓
┌─────────────────────────────────────────┐
│      Kafka Stream (partitioned)         │
│   - Partition by user_id for ordering   │
│   - Retention: 7 days for replay        │
└──────┬──────────────────────────────────┘
       ↓
┌─────────────────────────────────────────┐
│    Real-Time Feature Engineering        │
│  (Flink / Spark Streaming)              │
│                                         │
│  1. Velocity Features:                  │
│     - Transactions last 5/30/60 min     │
│     - Amount spent last 1 hour          │
│     - Unique merchants last 24h         │
│                                         │
│  2. Anomaly Features:                   │
│     - Unusual location (>500km from     │
│       last transaction)                 │
│     - New device fingerprint            │
│     - Unusual time (3am for daytime user)│
│                                         │
│  3. Network Features:                   │
│     - Merchant risk score               │
│     - IP reputation                     │
│     - Email/phone shared with fraudsters│
└──────┬──────────────────────────────────┘
       ↓
┌─────────────────────────────────────────┐
│         Feature Store Lookup            │
│   Online:  Redis (1-5ms)                │
│   Batch:   Cassandra/BigQuery           │
│                                         │
│   - User historical patterns            │
│   - Device fingerprints                 │
│   - Merchant metadata                   │
└──────┬──────────────────────────────────┘
       ↓
┌─────────────────────────────────────────┐
│         Multi-Layer Detection           │
│                                         │
│  Layer 1: Rule Engine (<10ms)           │
│   ├─ Blacklist check                    │
│   ├─ Amount thresholds                  │
│   └─ Basic velocity rules               │
│   → Block: 5% of fraud                  │
│                                         │
│  Layer 2: ML Model (<50ms)              │
│   ├─ Gradient Boosting (XGBoost)        │
│   ├─ Features: 200+                     │
│   └─ Score: 0-1 fraud probability       │
│   → Catch: 70% of fraud                 │
│                                         │
│  Layer 3: Deep Learning (<80ms)         │
│   ├─ LSTM for sequence modeling         │
│   ├─ Graph Neural Network               │
│   └─ Catches complex patterns           │
│   → Catch additional: 10% of fraud      │
└──────┬──────────────────────────────────┘
       ↓
┌─────────────────────────────────────────┐
│         Decision Logic                  │
│                                         │
│  if score > 0.9:                        │
│      → BLOCK (hard decline)             │
│  elif score > 0.7:                      │
│      → CHALLENGE (2FA, 3DS)             │
│  elif score > 0.5:                      │
│      → REVIEW (async manual review)     │
│  else:                                  │
│      → APPROVE                           │
└──────┬──────────────────────────────────┘
       ↓
┌─────────────────────────────────────────┐
│       Feedback Loop & Labeling          │
│   - User disputes (chargebacks)         │
│   - Manual review decisions             │
│   - Confirmed fraud cases               │
│   → Retrain models weekly               │
└─────────────────────────────────────────┘

class FraudDetectionSystem:
    def __init__(self):
        self.rule_engine = RuleEngine()
        self.ml_model = load_model('xgboost_v23.pkl')
        self.deep_model = load_model('lstm_v5.pt')
        self.feature_store = FeatureStore()
        self.decision_thresholds = {
            'block': 0.9,
            'challenge': 0.7,
            'review': 0.5
        }

    async def detect_fraud(self, transaction: dict) -> dict:
        start_time = time.time()

        # Step 1: Quick rule check (<5ms)
        rule_result = self.rule_engine.check(transaction)
        if rule_result['action'] == 'BLOCK':
            return {
                'decision': 'BLOCK',
                'reason': rule_result['reason'],
                'latency_ms': (time.time() - start_time) * 1000
            }

        # Step 2: Feature engineering (parallel)
        features = await asyncio.gather(
            self._compute_realtime_features(transaction),
            self._fetch_historical_features(transaction['user_id']),
            self._fetch_merchant_features(transaction['merchant_id'])
        )
        feature_vector = self._combine_features(*features)  # 200+ features

        # Step 3: ML scoring (<30ms)
        ml_score = self.ml_model.predict_proba(feature_vector)[0][1]

        # Step 4: Deep learning (only for borderline cases)
        if 0.4 < ml_score < 0.8:
            # Get transaction sequence for user
            sequence = await self._get_transaction_sequence(transaction['user_id'])
            dl_score = self.deep_model.predict(sequence)
            final_score = 0.6 * ml_score + 0.4 * dl_score
        else:
            final_score = ml_score

        # Step 5: Make decision
        decision = self._make_decision(final_score)

        # Step 6: Log for monitoring
        self._log_decision(transaction, final_score, decision)

        return {
            'decision': decision,
            'score': final_score,
            'latency_ms': (time.time() - start_time) * 1000
        }

    def _make_decision(self, score: float) -> str:
        if score > self.decision_thresholds['block']:
            return 'BLOCK'
        elif score > self.decision_thresholds['challenge']:
            return 'CHALLENGE'  # Ask for 2FA
        elif score > self.decision_thresholds['review']:
            return 'REVIEW'  # Manual review queue
        else:
            return 'APPROVE'

# Real-time Feature Engineering
class RealtimeFeatureEngine:
    def compute_velocity_features(self, user_id: str) -> dict:
        """Compute velocity over different time windows"""
        now = time.time()

        # Count transactions in time windows
        txns_5min = redis_client.zcount(f'txn:{user_id}', now - 300, now)
        txns_30min = redis_client.zcount(f'txn:{user_id}', now - 1800, now)
        txns_1hour = redis_client.zcount(f'txn:{user_id}', now - 3600, now)

        # Amount velocity
        amounts_1hour = redis_client.zrangebyscore(
            f'amt:{user_id}', now - 3600, now
        )
        total_amount_1hour = sum(float(a) for a in amounts_1hour)

        return {
            'txn_count_5min': txns_5min,
            'txn_count_30min': txns_30min,
            'txn_count_1hour': txns_1hour,
            'total_amount_1hour': total_amount_1hour,
            'avg_amount_1hour': total_amount_1hour / max(txns_1hour, 1)
        }

    def compute_anomaly_features(self, transaction: dict, user_profile: dict) -> dict:
        """Detect anomalies based on user history"""
        features = {}

        # Location anomaly
        last_location = user_profile.get('last_location')
        curr_location = (transaction['lat'], transaction['lon'])
        if last_location:
            distance_km = haversine_distance(last_location, curr_location)
            time_diff_hours = (transaction['timestamp'] - user_profile['last_txn_time']) / 3600
            features['distance_from_last'] = distance_km
            features['impossible_travel'] = 1 if distance_km > 1000 and time_diff_hours < 2 else 0

        # Amount anomaly (Z-score)
        avg_amount = user_profile.get('avg_transaction_amount', 100)
        std_amount = user_profile.get('std_transaction_amount', 50)
        features['amount_zscore'] = (transaction['amount'] - avg_amount) / std_amount

        # Time anomaly
        typical_hours = user_profile.get('typical_transaction_hours', [9, 10, 11, 14, 15, 16])
        current_hour = datetime.fromtimestamp(transaction['timestamp']).hour
        features['unusual_time'] = 1 if current_hour not in typical_hours else 0

        return features

# XGBoost Model (Primary)
model = xgb.XGBClassifier(
    n_estimators=500,
    max_depth=8,
    learning_rate=0.05,
    subsample=0.8,
    colsample_bytree=0.8,
    scale_pos_weight=10,  # Handle class imbalance (1:10 fraud:legit)
    eval_metric='auc'
)

# Features: 200+
feature_groups = {
    'transaction': 20,      # amount, merchant, category
    'velocity': 30,         # counts and amounts over time windows
    'anomaly': 15,          # deviations from user profile
    'network': 40,          # IP, device, email risk
    'behavioral': 50,       # spending patterns
    'merchant': 25,         # merchant risk, category
    'temporal': 20          # time-based features
}

# LSTM for Sequential Modeling
class FraudLSTM(nn.Module):
    def __init__(self):
        super().__init__()
        self.lstm = nn.LSTM(input_size=50, hidden_size=128, num_layers=2, batch_first=True)
        self.fc = nn.Sequential(
            nn.Linear(128, 64),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(64, 1),
            nn.Sigmoid()
        )

    def forward(self, sequence):
        # sequence: [batch, seq_len, 50 features]
        lstm_out, _ = self.lstm(sequence)
        last_hidden = lstm_out[:, -1, :]  # Take last timestep
        return self.fc(last_hidden)

# Key metrics to monitor
metrics = {
    'fraud_catch_rate': 0.80,  # Alert if drops below 75%
    'false_positive_rate': 0.01,  # Alert if exceeds 1.5%
    'p99_latency_ms': 100,  # Alert if exceeds 150ms
    'model_score_distribution': None,  # Alert on significant shift
    'feature_null_rate': 0.02,  # Alert if exceeds 5%
    'data_drift_psi': 0.15  # Alert if PSI > 0.25
}

┌──────────────────────────────────────────────────────┐
│              Feature Definition Layer                 │
│  - Python SDK for defining features                   │
│  - Schema validation and type checking                │
│  - Version control integration                        │
└──────────────┬───────────────────────────────────────┘
               ↓
┌──────────────────────────────────────────────────────┐
│           Feature Computation Layer                   │
│                                                       │
│  Batch (Spark/Dask):          Streaming (Flink):     │
│  - Daily aggregates            - Real-time counts     │
│  - Historical features         - Windowed aggregates  │
│  - Complex transformations     - Event-driven updates │
└──────────────┬───────────────────────────────────────┘
               ↓
┌──────────────────────────────────────────────────────┐
│              Feature Storage Layer                    │
│                                                       │
│  ┌────────────────────┐   ┌─────────────────────┐   │
│  │   Online Store     │   │   Offline Store     │   │
│  │  (Low Latency)     │   │   (Training Data)   │   │
│  │                    │   │                     │   │
│  │ Redis/DynamoDB     │   │ S3/BigQuery/Delta  │   │
│  │ - Key-value lookup │   │ - Point-in-time    │   │
│  │ - <5ms p99         │   │   joins            │   │
│  │ - Hot features     │   │ - Historical data  │   │
│  └────────────────────┘   └─────────────────────┘   │
└──────────────┬───────────────────────────────────────┘
               ↓
┌──────────────────────────────────────────────────────┐
│            Feature Registry (Metadata)                │
│  - Schema & types                                     │
│  - Lineage (data sources → features → models)        │
│  - Statistics (min, max, missing %)                  │
│  - Access control & governance                        │
└───────────────────────────────────────────────────────┘

from datetime import datetime, timedelta
from typing import List, Dict
import redis
import pandas as pd

# Feature Definition
class FeatureStore:
    def __init__(self):
        self.online_store = redis.Redis(host='localhost', port=6379)
        self.offline_store = BigQueryClient()
        self.registry = FeatureRegistry()

    # Define a feature
    @feature(
        name="user_purchase_count_7d",
        entity="user",
        value_type=ValueType.INT64,
        ttl=timedelta(days=7),
        online=True,
        offline=True
    )
    def user_purchase_count_7d(self, user_id: str, timestamp: datetime) -> int:
        """Count user purchases in last 7 days"""
        start_date = timestamp - timedelta(days=7)

        # For batch/training (point-in-time correct)
        if self.context == "offline":
            query = f"""
            SELECT user_id, COUNT(*) as purchase_count
            FROM purchases
            WHERE user_id = '{user_id}'
              AND purchase_timestamp >= '{start_date}'
              AND purchase_timestamp < '{timestamp}'
            GROUP BY user_id
            """
            return self.offline_store.query(query)

        # For online serving (real-time)
        else:
            # Pre-computed and cached in Redis
            key = f"user:{user_id}:purchase_count_7d"
            return int(self.online_store.get(key) or 0)

    # Get features for online serving
    def get_online_features(
        self,
        entity_rows: List[Dict],  # e.g., [{"user_id": "123"}, ...]
        feature_refs: List[str]   # e.g., ["user_purchase_count_7d", ...]
    ) -> pd.DataFrame:
        """
        Fast batch retrieval for inference
        Target latency: <5ms for 10 features
        """
        results = []

        # Parallel Redis MGET for performance
        pipeline = self.online_store.pipeline()

        for row in entity_rows:
            entity_key = f"user:{row['user_id']}"
            for feature in feature_refs:
                key = f"{entity_key}:{feature}"
                pipeline.get(key)

        # Execute all at once
        values = pipeline.execute()

        # Parse results
        idx = 0
        for row in entity_rows:
            feature_dict = {"user_id": row["user_id"]}
            for feature in feature_refs:
                feature_dict[feature] = values[idx]
                idx += 1
            results.append(feature_dict)

        return pd.DataFrame(results)

    # Get features for training (point-in-time correct)
    def get_historical_features(
        self,
        entity_df: pd.DataFrame,  # user_id, timestamp
        feature_refs: List[str]
    ) -> pd.DataFrame:
        """
        Point-in-time correct joins for training data
        Prevents data leakage
        """
        # Generate SQL with point-in-time joins
        query = self._build_pit_query(entity_df, feature_refs)

        # Execute on data warehouse
        result = self.offline_store.query(query)

        return result

    def _build_pit_query(self, entity_df, features):
        """
        Build SQL for point-in-time correct feature retrieval

        Example: If training data point is at 2024-01-15,
        only use features computed from data BEFORE 2024-01-15
        """
        base_query = """
        WITH entity_timestamps AS (
            SELECT user_id, event_timestamp
            FROM training_events
        )
        """

        # For each feature, join with timestamp constraint
        for feature in features:
            base_query += f"""
            LEFT JOIN LATERAL (
                SELECT {feature}
                FROM feature_values_{feature}
                WHERE entity_id = entity_timestamps.user_id
                  AND feature_timestamp <= entity_timestamps.event_timestamp
                ORDER BY feature_timestamp DESC
                LIMIT 1
            ) AS {feature}_values ON TRUE
            """

        return base_query

# Batch Feature Computation (Spark)
class BatchFeatureCompute:
    def compute_daily_features(self, date: datetime):
        """Run daily to compute batch features"""

        # Example: Compute user purchase count for all users
        query = """
        SELECT
            user_id,
            COUNT(*) as purchase_count_7d,
            SUM(amount) as total_spent_7d,
            AVG(amount) as avg_order_value_7d
        FROM purchases
        WHERE purchase_date BETWEEN {date - 7d} AND {date}
        GROUP BY user_id
        """

        df = spark.sql(query)

        # Write to both stores
        self._write_to_online_store(df)
        self._write_to_offline_store(df, date)

    def _write_to_online_store(self, df: DataFrame):
        """Write to Redis for low-latency serving"""
        # Batch write to Redis
        pipeline = redis_client.pipeline()

        for row in df.collect():
            key = f"user:{row.user_id}:purchase_count_7d"
            pipeline.set(key, row.purchase_count_7d, ex=7*24*3600)  # 7 day TTL

        pipeline.execute()

    def _write_to_offline_store(self, df: DataFrame, date: datetime):
        """Write to data warehouse for training"""
        # Append to partitioned table
        df.write.partitionBy("date").mode("append").saveAsTable(
            "feature_store.user_features"
        )

# Streaming Feature Computation (Flink)
class StreamingFeatureCompute:
    def process_realtime_event(self, event: dict):
        """Process events in real-time (Kafka → Flink → Redis)"""
        user_id = event['user_id']

        # Update velocity features
        current_count = redis_client.get(f"user:{user_id}:txn_count_1hr") or 0
        redis_client.incr(f"user:{user_id}:txn_count_1hr")
        redis_client.expire(f"user:{user_id}:txn_count_1hr", 3600)

        # Update windowed aggregates
        redis_client.zadd(
            f"user:{user_id}:recent_purchases",
            {event['purchase_id']: event['timestamp']}
        )

        # Remove old events outside window
        cutoff = time.time() - 3600
        redis_client.zremrangebyscore(
            f"user:{user_id}:recent_purchases",
            0,
            cutoff
        )

# WRONG: Data leakage - using future information
def get_features_WRONG(user_id, prediction_timestamp):
    # This query looks at ALL data, including future data!
    return db.query(f"""
        SELECT AVG(purchase_amount)
        FROM purchases
        WHERE user_id = '{user_id}'
    """)

# CORRECT: Point-in-time join
def get_features_CORRECT(user_id, prediction_timestamp):
    # Only use data from BEFORE prediction time
    return db.query(f"""
        SELECT AVG(purchase_amount)
        FROM purchases
        WHERE user_id = '{user_id}'
          AND purchase_timestamp < '{prediction_timestamp}'
    """)

# Key metrics
feature_metrics = {
    'online_latency_p99_ms': 5,
    'online_error_rate': 0.001,
    'feature_null_rate': {
        'user_purchase_count_7d': 0.02,  # 2% nulls acceptable
        'user_age': 0.10  # 10% nulls (optional feature)
    },
    'feature_staleness_minutes': {
        'batch_features': 24 * 60,  # Daily
        'streaming_features': 5      # 5 min max
    },
    'train_serve_skew': 0.05  # Feature distributions should match
}

┌─────────────────────────────────────────┐
│          API Gateway / Load Balancer    │
│  - Rate limiting (1000 QPS/user)        │
│  - Authentication & authorization        │
│  - Traffic routing by model version     │
└──────────────┬──────────────────────────┘
               ↓
┌─────────────────────────────────────────┐
│         Model Server Fleet              │
│  (Kubernetes pods with auto-scaling)    │
│                                         │
│  ┌────────────┐  ┌────────────┐       │
│  │  Server 1  │  │  Server 2  │  ...  │
│  │  CPU/GPU   │  │  CPU/GPU   │       │
│  │            │  │            │       │
│  │ Model A v1 │  │ Model A v2 │       │
│  │ Model B    │  │ Model C    │       │
│  └────────────┘  └────────────┘       │
└──────────────┬──────────────────────────┘
               ↓
┌─────────────────────────────────────────┐
│          Optimization Layer             │
│  - Request batching (collect 10-100ms)  │
│  - Result caching (Redis)               │
│  - Feature caching                      │
└──────────────┬──────────────────────────┘
               ↓
┌─────────────────────────────────────────┐
│         Model Registry & Storage        │
│  - S3/GCS: Model artifacts              │
│  - Versioning & metadata                │
│  - Lazy loading / preloading            │
└─────────────────────────────────────────┘

Monitoring:
┌─────────────────────────────────────────┐
│  Prometheus + Grafana + Alerts          │
│  - Latency (p50, p95, p99)              │
│  - Throughput (QPS)                     │
│  - GPU utilization                      │
│  - Model drift                          │
│  - Error rates                          │
└─────────────────────────────────────────┘

from fastapi import FastAPI, HTTPException
from typing import List, Dict
import torch
import numpy as np
import asyncio
from collections import defaultdict
import time

app = FastAPI()

class ModelServer:
    def __init__(self):
        self.models = {}  # model_name -> model
        self.batchers = {}  # model_name -> RequestBatcher
        self.cache = RedisCache()
        self.metrics = PrometheusMetrics()

    async def load_model(self, model_name: str, version: str):
        """Load model from registry"""
        # Download from S3/GCS
        model_path = f"s3://models/{model_name}/{version}/model.pt"

        if torch.cuda.is_available():
            device = torch.device("cuda")
            model = torch.load(model_path, map_location=device)
            model = torch.jit.script(model)  # TorchScript for optimization
        else:
            device = torch.device("cpu")
            model = torch.load(model_path, map_location=device)
            # Quantize for CPU inference
            model = torch.quantization.quantize_dynamic(
                model, {torch.nn.Linear}, dtype=torch.qint8
            )

        model.eval()
        self.models[model_name] = model
        self.batchers[model_name] = RequestBatcher(max_batch_size=32, max_wait_ms=50)

        print(f"Loaded {model_name} v{version} on {device}")

    @app.post("/predict/{model_name}")
    async def predict(self, model_name: str, features: Dict):
        """
        Prediction endpoint with batching and caching
        """
        start_time = time.time()

        # Step 1: Check cache
        cache_key = self._compute_cache_key(model_name, features)
        cached_result = await self.cache.get(cache_key)
        if cached_result:
            self.metrics.increment("cache_hit", model=model_name)
            return {"prediction": cached_result, "cached": True}

        # Step 2: Add to batch
        future = asyncio.Future()
        await self.batchers[model_name].add_request(features, future)

        # Wait for batch processing
        prediction = await future

        # Step 3: Cache result
        await self.cache.set(cache_key, prediction, ttl=3600)

        latency_ms = (time.time() - start_time) * 1000
        self.metrics.observe("prediction_latency", latency_ms, model=model_name)

        return {"prediction": prediction, "cached": False}

class RequestBatcher:
    """
    Batch requests for GPU efficiency
    Trade-off: Slight latency increase for much higher throughput
    """
    def __init__(self, max_batch_size=32, max_wait_ms=50):
        self.max_batch_size = max_batch_size
        self.max_wait_ms = max_wait_ms
        self.queue = []
        self.processing = False

    async def add_request(self, features, future):
        """Add request to batch queue"""
        self.queue.append((features, future))

        # Start batch processing if not already running
        if not self.processing:
            asyncio.create_task(self._process_batch())

        # Or if queue is full
        if len(self.queue) >= self.max_batch_size:
            asyncio.create_task(self._process_batch())

    async def _process_batch(self):
        """Process accumulated requests as batch"""
        if self.processing or len(self.queue) == 0:
            return

        self.processing = True

        # Wait for more requests (up to max_wait_ms)
        await asyncio.sleep(self.max_wait_ms / 1000)

        # Get batch
        batch_size = min(len(self.queue), self.max_batch_size)
        batch = self.queue[:batch_size]
        self.queue = self.queue[batch_size:]

        # Prepare batch tensor
        features_list = [item[0] for item in batch]
        futures = [item[1] for item in batch]

        # Convert to tensor
        batch_tensor = torch.tensor(
            np.array([self._features_to_array(f) for f in features_list])
        )

        # Run inference
        with torch.no_grad():
            predictions = model(batch_tensor)

        # Return results to individual futures
        for i, future in enumerate(futures):
            future.set_result(predictions[i].item())

        self.processing = False

        # Process remaining queue if any
        if len(self.queue) > 0:
            asyncio.create_task(self._process_batch())

# GPU Optimization
class GPUOptimizedServer:
    """Optimize for GPU serving"""

    def __init__(self):
        self.model = None
        self.use_amp = True  # Automatic Mixed Precision

    def load_optimized_model(self, model_path: str):
        """Load model with optimizations"""

        # TensorRT optimization (NVIDIA)
        import torch_tensorrt

        model = torch.load(model_path)

        # Compile with TensorRT
        trt_model = torch_tensorrt.compile(
            model,
            inputs=[torch_tensorrt.Input(shape=[1, 784])],
            enabled_precisions={torch.float16},  # FP16 for speed
            workspace_size=1 << 30  # 1GB
        )

        self.model = trt_model

    @torch.cuda.amp.autocast()  # Mixed precision
    def predict(self, batch_tensor):
        """Inference with AMP"""
        with torch.no_grad():
            return self.model(batch_tensor)

# A/B Testing Support
class ABTestingServer:
    """Route traffic to different model versions"""

    def __init__(self):
        self.model_versions = {
            'model_a': {'v1': 0.9, 'v2': 0.1},  # 90% v1, 10% v2
            'model_b': {'v1': 0.5, 'v2': 0.5}   # 50/50 split
        }

    def get_model_version(self, model_name: str, user_id: str) -> str:
        """Deterministic assignment based on user_id"""
        import hashlib

        # Hash user_id to get consistent assignment
        hash_value = int(hashlib.md5(user_id.encode()).hexdigest(), 16)
        bucket = (hash_value % 100) / 100.0

        # Assign to version based on bucket
        cumulative = 0
        for version, weight in self.model_versions[model_name].items():
            cumulative += weight
            if bucket < cumulative:
                return version

        return 'v1'  # Default

# Auto-scaling based on metrics
class AutoScaler:
    """Scale model servers based on load"""

    def should_scale_up(self, metrics):
        """Decide if we need more servers"""
        conditions = [
            metrics['cpu_usage'] > 80,
            metrics['gpu_usage'] > 85,
            metrics['p99_latency_ms'] > 100,
            metrics['queue_size'] > 1000
        ]

        return any(conditions)

    def should_scale_down(self, metrics):
        """Decide if we can reduce servers"""
        conditions = [
            metrics['cpu_usage'] < 30,
            metrics['gpu_usage'] < 30,
            metrics['p99_latency_ms'] < 20,
            metrics['queue_size'] < 100
        ]

        return all(conditions)

# Key metrics to track
serving_metrics = {
    'latency_p50_ms': 10,
    'latency_p95_ms': 30,
    'latency_p99_ms': 50,
    'qps': 10000,
    'error_rate': 0.001,
    'gpu_utilization_%': 75,
    'gpu_memory_used_gb': 10,
    'batch_size_avg': 24,
    'cache_hit_rate': 0.30,
    'model_load_time_s': 15
}

# Alerts
alerts = {
    'p99_latency > 100ms': 'High latency',
    'error_rate > 0.01': 'High error rate',
    'gpu_util < 40%': 'Underutilized GPU',
    'qps drops > 50%': 'Traffic drop'
}

┌────────────────────────────────────────────────┐
│         Production Predictions                 │
│  (Model serving logs every prediction)         │
└──────────────┬─────────────────────────────────┘
               ↓
┌────────────────────────────────────────────────┐
│          Streaming Pipeline (Kafka)            │
│  - Prediction logs                             │
│  - Features used                               │
│  - Model version                               │
│  - Latency, errors                             │
└──────────────┬─────────────────────────────────┘
               ↓
┌────────────────────────────────────────────────┐
│         Real-Time Monitoring Layer             │
│  (Flink/Spark Streaming)                       │
│                                                │
│  1. Data Quality Checks:                       │
│     - Schema validation                        │
│     - Missing value detection                  │
│     - Range/distribution checks                │
│                                                │
│  2. Data Drift Detection:                      │
│     - PSI (Population Stability Index)         │
│     - KL Divergence                            │
│     - Kolmogorov-Smirnov test                  │
│                                                │
│  3. Performance Monitoring:                    │
│     - Latency (p50, p95, p99)                  │
│     - Throughput (QPS)                         │
│     - Error rates                              │
└──────────────┬─────────────────────────────────┘
               ↓
┌────────────────────────────────────────────────┐
│         Ground Truth Collection                │
│  (Delayed labels via user feedback)            │
│  - User clicks/conversions                     │
│  - Manual labels                               │
│  - Downstream outcomes                         │
└──────────────┬─────────────────────────────────┘
               ↓
┌────────────────────────────────────────────────┐
│         Model Performance Analysis             │
│  (Daily batch jobs)                            │
│                                                │
│  - Accuracy, Precision, Recall                 │
│  - AUC, F1 score                               │
│  - Per-segment performance                     │
│  - Calibration metrics                         │
└──────────────┬─────────────────────────────────┘
               ↓
┌────────────────────────────────────────────────┐
│        Alerting & Visualization                │
│                                                │
│  - Prometheus + Grafana dashboards             │
│  - PagerDuty alerts                            │
│  - Weekly performance reports                  │
└────────────────────────────────────────────────┘

import numpy as np
from scipy import stats
from typing import Dict, List
import pandas as pd

class ModelMonitor:
    def __init__(self, model_name: str):
        self.model_name = model_name
        self.baseline_stats = self._load_baseline_stats()
        self.alert_thresholds = {
            'psi': 0.2,
            'kl_divergence': 0.1,
            'accuracy_drop': 0.05,
            'p99_latency_ms': 100,
            'error_rate': 0.01
        }

    # 1. DATA QUALITY MONITORING
    def check_data_quality(self, batch: pd.DataFrame) -> Dict:
        """Real-time data quality checks"""
        issues = []

        # Schema validation
        expected_cols = set(self.baseline_stats['feature_names'])
        actual_cols = set(batch.columns)
        if expected_cols != actual_cols:
            issues.append({
                'type': 'SCHEMA_DRIFT',
                'severity': 'CRITICAL',
                'message': f'Missing columns: {expected_cols - actual_cols}'
            })

        # Missing values
        missing_pct = batch.isnull().sum() / len(batch)
        high_missing = missing_pct[missing_pct > 0.1]
        if len(high_missing) > 0:
            issues.append({
                'type': 'HIGH_MISSING_VALUES',
                'severity': 'WARNING',
                'features': high_missing.to_dict()
            })

        # Range validation
        for col in batch.select_dtypes(include=[np.number]).columns:
            baseline_min = self.baseline_stats['ranges'][col]['min']
            baseline_max = self.baseline_stats['ranges'][col]['max']

            current_min = batch[col].min()
            current_max = batch[col].max()

            if current_min < baseline_min * 0.5 or current_max > baseline_max * 2:
                issues.append({
                    'type': 'OUT_OF_RANGE',
                    'severity': 'WARNING',
                    'feature': col,
                    'baseline': f'[{baseline_min}, {baseline_max}]',
                    'current': f'[{current_min}, {current_max}]'
                })

        return {'issues': issues, 'passed': len(issues) == 0}

    # 2. DATA DRIFT DETECTION
    def detect_data_drift(self, current_data: pd.DataFrame) -> Dict:
        """Detect feature distribution drift"""
        drift_results = {}

        for feature in current_data.columns:
            if feature in self.baseline_stats['distributions']:
                # PSI (Population Stability Index)
                psi = self._calculate_psi(
                    self.baseline_stats['distributions'][feature],
                    current_data[feature]
                )

                # KL Divergence
                kl_div = self._calculate_kl_divergence(
                    self.baseline_stats['distributions'][feature],
                    current_data[feature]
                )

                # Kolmogorov-Smirnov test
                ks_stat, ks_pvalue = stats.ks_2samp(
                    self.baseline_stats['distributions'][feature],
                    current_data[feature]
                )

                drift_results[feature] = {
                    'psi': psi,
                    'kl_divergence': kl_div,
                    'ks_statistic': ks_stat,
                    'ks_pvalue': ks_pvalue,
                    'drifted': psi > self.alert_thresholds['psi']
                }

        return drift_results

    def _calculate_psi(self, baseline: np.ndarray, current: np.ndarray, bins=10) -> float:
        """
        Population Stability Index
        PSI < 0.1: No significant drift
        0.1 < PSI < 0.2: Moderate drift
        PSI > 0.2: Significant drift
        """
        # Create bins from baseline
        breakpoints = np.percentile(baseline, np.linspace(0, 100, bins + 1))
        breakpoints[-1] += 0.001  # Include max value

        # Calculate distributions
        baseline_counts = np.histogram(baseline, bins=breakpoints)[0]
        current_counts = np.histogram(current, bins=breakpoints)[0]

        # Convert to percentages
        baseline_pct = baseline_counts / len(baseline)
        current_pct = current_counts / len(current)

        # Avoid division by zero
        baseline_pct = np.where(baseline_pct == 0, 0.0001, baseline_pct)
        current_pct = np.where(current_pct == 0, 0.0001, current_pct)

        # PSI formula
        psi = np.sum((current_pct - baseline_pct) * np.log(current_pct / baseline_pct))
        return psi

    def _calculate_kl_divergence(self, baseline: np.ndarray, current: np.ndarray, bins=50) -> float:
        """KL Divergence: D_KL(P||Q)"""
        # Create histograms
        hist_range = (min(baseline.min(), current.min()),
                     max(baseline.max(), current.max()))

        p, _ = np.histogram(baseline, bins=bins, range=hist_range, density=True)
        q, _ = np.histogram(current, bins=bins, range=hist_range, density=True)

        # Normalize and avoid zeros
        p = p / p.sum()
        q = q / q.sum()
        p = np.where(p == 0, 1e-10, p)
        q = np.where(q == 0, 1e-10, q)

        # KL divergence
        kl = np.sum(p * np.log(p / q))
        return kl

    # 3. MODEL PERFORMANCE MONITORING
    def monitor_model_performance(
        self,
        predictions: np.ndarray,
        actuals: np.ndarray,
        prediction_times: List[float]
    ) -> Dict:
        """Monitor model accuracy and performance"""
        from sklearn.metrics import accuracy_score, roc_auc_score, precision_recall_fscore_support

        metrics = {}

        # Classification metrics (if labels available)
        if actuals is not None:
            metrics['accuracy'] = accuracy_score(actuals, predictions > 0.5)
            metrics['auc'] = roc_auc_score(actuals, predictions)

            precision, recall, f1, _ = precision_recall_fscore_support(
                actuals, predictions > 0.5, average='binary'
            )
            metrics['precision'] = precision
            metrics['recall'] = recall
            metrics['f1'] = f1

            # Check for degradation
            baseline_accuracy = self.baseline_stats['accuracy']
            if metrics['accuracy'] < baseline_accuracy - self.alert_thresholds['accuracy_drop']:
                self._trigger_alert({
                    'type': 'ACCURACY_DROP',
                    'severity': 'CRITICAL',
                    'baseline': baseline_accuracy,
                    'current': metrics['accuracy'],
                    'drop': baseline_accuracy - metrics['accuracy']
                })

        # Latency monitoring
        latency_p50 = np.percentile(prediction_times, 50)
        latency_p95 = np.percentile(prediction_times, 95)
        latency_p99 = np.percentile(prediction_times, 99)

        metrics['latency_ms'] = {
            'p50': latency_p50,
            'p95': latency_p95,
            'p99': latency_p99
        }

        if latency_p99 > self.alert_thresholds['p99_latency_ms']:
            self._trigger_alert({
                'type': 'HIGH_LATENCY',
                'severity': 'WARNING',
                'p99_latency': latency_p99,
                'threshold': self.alert_thresholds['p99_latency_ms']
            })

        return metrics

    # 4. PREDICTION DRIFT (MODEL OUTPUT DISTRIBUTION)
    def monitor_prediction_drift(self, predictions: np.ndarray) -> Dict:
        """Check if prediction distribution has changed"""
        # For classification: check score distribution
        score_buckets = [0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0]
        current_dist = np.histogram(predictions, bins=score_buckets)[0]
        current_dist = current_dist / current_dist.sum()

        baseline_dist = self.baseline_stats['prediction_distribution']

        # Chi-square test
        chi_stat, p_value = stats.chisquare(current_dist, baseline_dist)

        return {
            'chi_square_statistic': chi_stat,
            'p_value': p_value,
            'drifted': p_value < 0.05,  # Significant at 5% level
            'current_distribution': current_dist.tolist(),
            'baseline_distribution': baseline_dist.tolist()
        }

    # 5. BUSINESS METRICS MONITORING
    def monitor_business_metrics(self, predictions: pd.DataFrame, outcomes: pd.DataFrame) -> Dict:
        """Monitor business impact"""
        # Example: For a recommendation system
        metrics = {
            'ctr': outcomes['clicked'].mean(),
            'conversion_rate': outcomes['converted'].mean(),
            'revenue_per_impression': outcomes['revenue'].mean(),
            'engagement_time': outcomes['time_spent'].mean()
        }

        # Compare with baseline
        for metric, value in metrics.items():
            baseline = self.baseline_stats['business_metrics'][metric]
            change_pct = (value - baseline) / baseline * 100

            if abs(change_pct) > 10:  # 10% change threshold
                self._trigger_alert({
                    'type': 'BUSINESS_METRIC_CHANGE',
                    'severity': 'WARNING',
                    'metric': metric,
                    'baseline': baseline,
                    'current': value,
                    'change_pct': change_pct
                })

        return metrics

    def _trigger_alert(self, alert: Dict):
        """Send alert to monitoring system"""
        print(f"🚨 ALERT: {alert['type']} - {alert['severity']}")
        # Send to PagerDuty, Slack, etc.
        self._send_to_pagerduty(alert)
        self._send_to_slack(alert)

drift_severity = {
    'psi': {
        'low': (0, 0.1),      # No action needed
        'medium': (0.1, 0.2),  # Investigate
        'high': (0.2, float('inf'))  # Retrain model
    },
    'kl_divergence': {
        'low': (0, 0.05),
        'medium': (0.05, 0.1),
        'high': (0.1, float('inf'))
    }
}

class AutoRemediation:
    def handle_drift(self, drift_severity: str):
        """Automated response to drift"""
        if drift_severity == 'high':
            # Trigger model retraining
            self.trigger_retraining_pipeline()

            # Meanwhile, rollback to previous version
            self.rollback_model_version()

        elif drift_severity == 'medium':
            # Increase monitoring frequency
            self.increase_monitoring_frequency()

            # Alert data science team
            self.alert_team()

┌──────────────────────────────────────────────────────┐
│             Orchestration Layer                       │
│  Kubernetes + Kubeflow / Ray / Slurm                  │
│  - Resource allocation                                │
│  - Fault tolerance & checkpointing                    │
│  - Job scheduling                                     │
└──────────────┬───────────────────────────────────────┘
               ↓
┌──────────────────────────────────────────────────────┐
│          Data Parallelism (Most Common)               │
│                                                       │
│  GPU 1: Model copy 1 → Batch 1 → Gradients           │
│  GPU 2: Model copy 2 → Batch 2 → Gradients           │
│  GPU 3: Model copy 3 → Batch 3 → Gradients           │
│  GPU 4: Model copy 4 → Batch 4 → Gradients           │
│                         ↓                             │
│              All-Reduce (Average gradients)           │
│                         ↓                             │
│              Update all model copies                  │
└───────────────────────────────────────────────────────┘

For VERY large models:
┌──────────────────────────────────────────────────────┐
│          Model Parallelism (Layers split)             │
│                                                       │
│  GPU 1: Layers 1-25    → Forward → Activation        │
│  GPU 2: Layers 26-50   → Forward → Activation        │
│  GPU 3: Layers 51-75   → Forward → Activation        │
│  GPU 4: Layers 76-100  → Forward → Output            │
│                                                       │
│  Backward pass flows in reverse                       │
└───────────────────────────────────────────────────────┘

┌──────────────────────────────────────────────────────┐
│       Pipeline Parallelism (Micro-batching)           │
│                                                       │
│  Time  │  GPU 1  │  GPU 2  │  GPU 3  │  GPU 4       │
│  ──────┼─────────┼─────────┼─────────┼──────────    │
│   t1   │ Batch 1 │    -    │    -    │    -         │
│   t2   │ Batch 2 │ Batch 1 │    -    │    -         │
│   t3   │ Batch 3 │ Batch 2 │ Batch 1 │    -         │
│   t4   │ Batch 4 │ Batch 3 │ Batch 2 │ Batch 1      │
│                                                       │
│  Minimize bubble (idle time) with micro-batches       │
└───────────────────────────────────────────────────────┘

import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data.distributed import DistributedSampler

# 1. DATA PARALLEL (Most Common) - PyTorch
def setup_distributed():
    """Initialize distributed training"""
    # Initialize process group
    dist.init_process_group(
        backend='nccl',  # NVIDIA Collective Communications Library
        init_method='env://',  # Use environment variables
        world_size=int(os.environ['WORLD_SIZE']),  # Total GPUs
        rank=int(os.environ['RANK'])  # This GPU's rank
    )

def train_data_parallel(model, train_dataset, epochs=10):
    """Data parallel training"""
    # Setup
    setup_distributed()
    local_rank = int(os.environ['LOCAL_RANK'])
    device = torch.device(f'cuda:{local_rank}')

    # Wrap model with DDP
    model = model.to(device)
    model = DDP(model, device_ids=[local_rank])

    # Distributed sampler (each GPU gets different data)
    sampler = DistributedSampler(
        train_dataset,
        num_replicas=dist.get_world_size(),
        rank=dist.get_rank(),
        shuffle=True
    )

    dataloader = DataLoader(
        train_dataset,
        batch_size=32,
        sampler=sampler,
        num_workers=4,
        pin_memory=True
    )

    optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)

    for epoch in range(epochs):
        # Set epoch for shuffling
        sampler.set_epoch(epoch)

        for batch_idx, (data, target) in enumerate(dataloader):
            data, target = data.to(device), target.to(device)

            # Forward pass
            output = model(data)
            loss = F.cross_entropy(output, target)

            # Backward pass
            optimizer.zero_grad()
            loss.backward()  # Gradients are automatically all-reduced by DDP

            # Update weights
            optimizer.step()

            # Logging (only rank 0)
            if dist.get_rank() == 0 and batch_idx % 100 == 0:
                print(f'Epoch {epoch}, Loss: {loss.item():.4f}')

        # Checkpoint (only rank 0)
        if dist.get_rank() == 0:
            torch.save({
                'epoch': epoch,
                'model_state_dict': model.module.state_dict(),
                'optimizer_state_dict': optimizer.state_dict(),
            }, f'checkpoint_epoch_{epoch}.pt')

# 2. MODEL PARALLEL - For Large Models
class ModelParallelTransformer(nn.Module):
    """Split large model across GPUs"""
    def __init__(self, num_layers=96, hidden_size=12288):
        super().__init__()

        # Split layers across 4 GPUs
        layers_per_gpu = num_layers // 4

        # GPU 0: First 25% of layers
        self.layers_0 = nn.Sequential(*[
            TransformerBlock(hidden_size) for _ in range(layers_per_gpu)
        ]).to('cuda:0')

        # GPU 1: Next 25%
        self.layers_1 = nn.Sequential(*[
            TransformerBlock(hidden_size) for _ in range(layers_per_gpu)
        ]).to('cuda:1')

        # GPU 2: Next 25%
        self.layers_2 = nn.Sequential(*[
            TransformerBlock(hidden_size) for _ in range(layers_per_gpu)
        ]).to('cuda:2')

        # GPU 3: Last 25%
        self.layers_3 = nn.Sequential(*[
            TransformerBlock(hidden_size) for _ in range(layers_per_gpu)
        ]).to('cuda:3')

        self.output = nn.Linear(hidden_size, vocab_size).to('cuda:3')

    def forward(self, x):
        # Move through GPUs sequentially
        x = x.to('cuda:0')
        x = self.layers_0(x)

        x = x.to('cuda:1')
        x = self.layers_1(x)

        x = x.to('cuda:2')
        x = self.layers_2(x)

        x = x.to('cuda:3')
        x = self.layers_3(x)
        x = self.output(x)

        return x

# 3. PIPELINE PARALLEL - Deepspeed, Megatron-LM
from deepspeed.pipe import PipelineModule, LayerSpec

def pipeline_parallel():
    """Pipeline parallelism with DeepSpeed"""
    # Define model as sequence of layers
    layers = [
        LayerSpec(TransformerBlock, args=(hidden_size,))
        for _ in range(96)
    ]

    # DeepSpeed will automatically partition across GPUs
    model = PipelineModule(
        layers=layers,
        num_stages=4,  # 4 GPUs
        partition_method='uniform'  # or 'balanced'
    )

    # Training with micro-batches
    engine, _, _, _ = deepspeed.initialize(
        model=model,
        config={
            'train_micro_batch_size_per_gpu': 4,
            'gradient_accumulation_steps': 4,
            'pipeline': {
                'pipe_partitioned': True,
                'grad_partitioned': True
            }
        }
    )

    for batch in dataloader:
        loss = engine(batch)
        engine.backward(loss)
        engine.step()

# 4. ZERO OPTIMIZER (Memory Optimization)
from deepspeed import DeepSpeedConfig

def train_with_zero():
    """ZeRO: Memory-optimized distributed training"""
    # ZeRO Stage 1: Partition optimizer states
    # ZeRO Stage 2: + Partition gradients
    # ZeRO Stage 3: + Partition model parameters

    config = {
        "train_batch_size": 128,
        "gradient_accumulation_steps": 4,
        "zero_optimization": {
            "stage": 3,  # Full ZeRO
            "offload_optimizer": {
                "device": "cpu",  # Offload to CPU RAM
                "pin_memory": True
            },
            "offload_param": {
                "device": "cpu"
            }
        },
        "fp16": {
            "enabled": True  # Mixed precision
        }
    }

    model_engine, optimizer, _, _ = deepspeed.initialize(
        model=model,
        model_parameters=model.parameters(),
        config=config
    )

# 5. GRADIENT ACCUMULATION (Simulate larger batch)
def train_with_gradient_accumulation(model, dataloader, accumulation_steps=4):
    """Accumulate gradients before update"""
    optimizer.zero_grad()

    for i, (data, target) in enumerate(dataloader):
        output = model(data)
        loss = criterion(output, target)

        # Scale loss by accumulation steps
        loss = loss / accumulation_steps
        loss.backward()

        # Update every N steps
        if (i + 1) % accumulation_steps == 0:
            optimizer.step()
            optimizer.zero_grad()

# 1. Mixed Precision Training (FP16)
from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()

for data, target in dataloader:
    optimizer.zero_grad()

    # Forward in FP16
    with autocast():
        output = model(data)
        loss = criterion(output, target)

    # Backward with gradient scaling
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

# 2. Gradient Checkpointing (Memory Savings)
from torch.utils.checkpoint import checkpoint

class CheckpointedBlock(nn.Module):
    def forward(self, x):
        # Don't store activations, recompute in backward
        return checkpoint(self._forward, x)

    def _forward(self, x):
        return self.layer(x)

# 3. Gradient Compression
class GradientCompressor:
    def compress(self, tensor, compression_ratio=0.01):
        """Top-k gradient sparsification"""
        numel = tensor.numel()
        k = max(1, int(numel * compression_ratio))

        # Keep only top-k gradients
        values, indices = torch.topk(tensor.abs().flatten(), k)
        compressed = torch.zeros_like(tensor.flatten())
        compressed[indices] = tensor.flatten()[indices]

        return compressed.reshape(tensor.shape)

class FaultTolerantTrainer:
    def __init__(self, checkpoint_freq=100):
        self.checkpoint_freq = checkpoint_freq

    def save_checkpoint(self, epoch, model, optimizer, path):
        """Save training state"""
        if dist.get_rank() == 0:  # Only rank 0 saves
            torch.save({
                'epoch': epoch,
                'model_state_dict': model.state_dict(),
                'optimizer_state_dict': optimizer.state_dict(),
                'rng_state': torch.get_rng_state(),
                'cuda_rng_state': torch.cuda.get_rng_state_all()
            }, path)

    def load_checkpoint(self, path, model, optimizer):
        """Resume from checkpoint"""
        checkpoint = torch.load(path)
        model.load_state_dict(checkpoint['model_state_dict'])
        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
        torch.set_rng_state(checkpoint['rng_state'])
        torch.cuda.set_rng_state_all(checkpoint['cuda_rng_state'])
        return checkpoint['epoch']

┌────────────────────────────────────────────────────┐
│           Experiment Configuration                 │
│  - Define variants (A, B, C)                       │
│  - Traffic allocation (50/50, 90/10, etc.)         │
│  - Target audience (location, platform, etc.)      │
│  - Metrics (primary, secondary, guardrails)        │
│  - Duration & sample size calculation              │
└──────────────┬─────────────────────────────────────┘
               ↓
┌────────────────────────────────────────────────────┐
│        Assignment Service (User Bucketing)         │
│                                                    │
│  Input: user_id, experiment_id                     │
│  Output: variant (A or B)                          │
│                                                    │
│  hash(user_id + experiment_id) % 100               │
│    → Deterministic, consistent assignment          │
│    → Same user always gets same variant            │
│    → No database lookup needed                     │
└──────────────┬─────────────────────────────────────┘
               ↓
┌────────────────────────────────────────────────────┐
│         User Experience (Application)              │
│                                                    │
│  if variant == 'A':                                │
│      show_old_checkout_flow()                      │
│  elif variant == 'B':                              │
│      show_new_checkout_flow()                      │
└──────────────┬─────────────────────────────────────┘
               ↓
┌────────────────────────────────────────────────────┐
│          Event Tracking (Kafka Stream)             │
│                                                    │
│  - Exposure events (user saw variant)              │
│  - Action events (clicks, purchases, etc.)         │
│  - Metadata (timestamp, user_id, variant, etc.)   │
└──────────────┬─────────────────────────────────────┘
               ↓
┌────────────────────────────────────────────────────┐
│       Data Pipeline (Batch Processing)             │
│                                                    │
│  Daily Spark jobs:                                 │
│  - Join exposure + outcome events                  │
│  - Calculate metrics per variant                   │
│  - Run statistical tests                           │
│  - Detect Sample Ratio Mismatch (SRM)             │
└──────────────┬─────────────────────────────────────┘
               ↓
┌────────────────────────────────────────────────────┐
│         Statistical Analysis Engine                │
│                                                    │
│  - T-test for continuous metrics                   │
│  - Z-test for proportions                          │
│  - Sequential testing (early stopping)             │
│  - Multiple testing correction (Bonferroni)        │
│  - Variance reduction (CUPED, stratification)      │
└──────────────┬─────────────────────────────────────┘
               ↓
┌────────────────────────────────────────────────────┐
│      Dashboard & Reporting (Real-time)             │
│                                                    │
│  - Experiment status & health                      │
│  - Metric movements (% change, confidence)         │
│  - Statistical significance & p-values             │
│  - Sample Ratio Mismatch alerts                    │
│  - Interaction effects detection                   │
└────────────────────────────────────────────────────┘

import hashlib
import numpy as np
from scipy import stats
from typing import Dict, List, Tuple

# 1. ASSIGNMENT SERVICE
class ExperimentAssignmentService:
    """Deterministic user assignment to experiment variants"""

    def __init__(self):
        self.experiments = {}  # experiment_id -> config

    def assign_variant(self, user_id: str, experiment_id: str) -> str:
        """
        Deterministic assignment using hash function
        Same user always gets same variant
        """
        experiment = self.experiments[experiment_id]

        # Hash user_id + experiment_id for randomization
        hash_input = f"{user_id}_{experiment_id}"
        hash_value = int(hashlib.md5(hash_input.encode()).hexdigest(), 16)

        # Convert to bucket (0-99)
        bucket = hash_value % 100

        # Assign to variant based on traffic allocation
        cumulative = 0
        for variant, allocation in experiment['traffic_allocation'].items():
            cumulative += allocation
            if bucket < cumulative:
                return variant

        return 'control'  # Default

    def should_include_user(
        self,
        user: Dict,
        experiment_config: Dict
    ) -> bool:
        """Check if user qualifies for experiment"""
        targeting = experiment_config['targeting']

        # Check filters
        if 'countries' in targeting:
            if user['country'] not in targeting['countries']:
                return False

        if 'platforms' in targeting:
            if user['platform'] not in targeting['platforms']:
                return False

        if 'user_segments' in targeting:
            if user['segment'] not in targeting['user_segments']:
                return False

        return True

# 2. EVENT TRACKING
class ExperimentEventTracker:
    """Track exposure and outcome events"""

    def track_exposure(
        self,
        user_id: str,
        experiment_id: str,
        variant: str,
        timestamp: int
    ):
        """Log when user is exposed to experiment"""
        event = {
            'event_type': 'exposure',
            'user_id': user_id,
            'experiment_id': experiment_id,
            'variant': variant,
            'timestamp': timestamp
        }
        self._send_to_kafka('experiment_events', event)

    def track_outcome(
        self,
        user_id: str,
        experiment_id: str,
        metric_name: str,
        metric_value: float,
        timestamp: int
    ):
        """Log outcome metric (conversion, revenue, etc.)"""
        event = {
            'event_type': 'outcome',
            'user_id': user_id,
            'experiment_id': experiment_id,
            'metric_name': metric_name,
            'metric_value': metric_value,
            'timestamp': timestamp
        }
        self._send_to_kafka('experiment_events', event)

# 3. STATISTICAL ANALYSIS
class ExperimentAnalyzer:
    """Analyze experiment results"""

    def __init__(self):
        self.alpha = 0.05  # Significance level (5%)
        self.power = 0.80  # Statistical power (80%)

    def calculate_sample_size(
        self,
        baseline_rate: float,
        minimum_detectable_effect: float,
        alpha: float = 0.05,
        power: float = 0.80
    ) -> int:
        """
        Calculate required sample size per variant
        For detecting a minimum effect with desired power
        """
        from statsmodels.stats.power import zt_ind_solve_power

        # Effect size (Cohen's h for proportions)
        p1 = baseline_rate
        p2 = baseline_rate * (1 + minimum_detectable_effect)

        effect_size = 2 * (np.arcsin(np.sqrt(p2)) - np.arcsin(np.sqrt(p1)))

        # Calculate sample size
        n = zt_ind_solve_power(
            effect_size=effect_size,
            alpha=alpha,
            power=power,
            alternative='two-sided'
        )

        return int(np.ceil(n))

    def analyze_experiment(
        self,
        control_metrics: np.ndarray,
        treatment_metrics: np.ndarray
    ) -> Dict:
        """
        Run statistical test on experiment results
        """
        n_control = len(control_metrics)
        n_treatment = len(treatment_metrics)

        mean_control = np.mean(control_metrics)
        mean_treatment = np.mean(treatment_metrics)

        # Relative lift
        relative_lift = (mean_treatment - mean_control) / mean_control

        # T-test for continuous metrics
        t_stat, p_value = stats.ttest_ind(
            treatment_metrics,
            control_metrics,
            equal_var=False  # Welch's t-test
        )

        # Confidence interval (95%)
        se_diff = np.sqrt(
            np.var(control_metrics) / n_control +
            np.var(treatment_metrics) / n_treatment
        )
        ci_lower = (mean_treatment - mean_control) - 1.96 * se_diff
        ci_upper = (mean_treatment - mean_control) + 1.96 * se_diff

        is_significant = p_value < self.alpha

        return {
            'control_mean': mean_control,
            'treatment_mean': mean_treatment,
            'absolute_lift': mean_treatment - mean_control,
            'relative_lift': relative_lift,
            'p_value': p_value,
            'is_significant': is_significant,
            'confidence_interval': (ci_lower, ci_upper),
            'sample_size_control': n_control,
            'sample_size_treatment': n_treatment
        }

    def check_sample_ratio_mismatch(
        self,
        n_control: int,
        n_treatment: int,
        expected_ratio: float = 0.5
    ) -> Dict:
        """
        Sample Ratio Mismatch (SRM) detection
        Checks if traffic split matches expected ratio
        """
        total = n_control + n_treatment
        expected_control = total * expected_ratio
        expected_treatment = total * (1 - expected_ratio)

        # Chi-square test
        observed = [n_control, n_treatment]
        expected = [expected_control, expected_treatment]

        chi_stat, p_value = stats.chisquare(observed, expected)

        has_srm = p_value < 0.001  # Very strict threshold

        return {
            'n_control': n_control,
            'n_treatment': n_treatment,
            'expected_ratio': expected_ratio,
            'actual_ratio': n_control / total,
            'p_value': p_value,
            'has_srm': has_srm
        }

    def apply_cuped(
        self,
        post_metrics: np.ndarray,
        pre_metrics: np.ndarray
    ) -> np.ndarray:
        """
        CUPED (Controlled-experiment Using Pre-Experiment Data)
        Variance reduction technique using covariates
        """
        # Calculate theta (optimal coefficient)
        cov = np.cov(post_metrics, pre_metrics)[0, 1]
        var_pre = np.var(pre_metrics)
        theta = cov / var_pre

        # Adjust post-experiment metric
        adjusted_metrics = post_metrics - theta * (pre_metrics - np.mean(pre_metrics))

        # Variance reduction
        var_original = np.var(post_metrics)
        var_adjusted = np.var(adjusted_metrics)
        variance_reduction = 1 - (var_adjusted / var_original)

        print(f"Variance reduced by {variance_reduction:.1%}")

        return adjusted_metrics

    def sequential_testing(
        self,
        control_data: List[float],
        treatment_data: List[float],
        looks: int = 5
    ) -> Dict:
        """
        Sequential testing for early stopping
        Allows peeking at results without inflating false positive rate
        """
        # Always-valid p-values (mixture sequential probability ratio test)
        results = []

        for i in range(1, looks + 1):
            # Get data up to this point
            idx = int(len(control_data) * i / looks)
            control_subset = control_data[:idx]
            treatment_subset = treatment_data[:idx]

            # Run test
            result = self.analyze_experiment(
                np.array(control_subset),
                np.array(treatment_subset)
            )

            # Adjusted alpha for multiple looks (Bonferroni correction)
            adjusted_alpha = self.alpha / looks
            result['adjusted_alpha'] = adjusted_alpha
            result['can_stop'] = result['p_value'] < adjusted_alpha

            results.append(result)

            if result['can_stop']:
                print(f"Can stop early at look {i}/{looks}")
                break

        return results

# 4. INTERACTION EFFECTS
class InteractionEffectsDetector:
    """Detect when multiple experiments interfere"""

    def detect_interaction(
        self,
        exp1_assignment: np.ndarray,  # 0 or 1
        exp2_assignment: np.ndarray,  # 0 or 1
        outcome: np.ndarray
    ) -> Dict:
        """
        2-way ANOVA to detect interaction effects
        """
        from scipy.stats import f_oneway

        # Four groups: (exp1=0, exp2=0), (exp1=1, exp2=0), etc.
        group_00 = outcome[(exp1_assignment == 0) & (exp2_assignment == 0)]
        group_01 = outcome[(exp1_assignment == 0) & (exp2_assignment == 1)]
        group_10 = outcome[(exp1_assignment == 1) & (exp2_assignment == 0)]
        group_11 = outcome[(exp1_assignment == 1) & (exp2_assignment == 1)]

        # Main effect of exp1
        exp1_control = np.concatenate([group_00, group_01])
        exp1_treatment = np.concatenate([group_10, group_11])
        _, p_exp1 = stats.ttest_ind(exp1_treatment, exp1_control)

        # Main effect of exp2
        exp2_control = np.concatenate([group_00, group_10])
        exp2_treatment = np.concatenate([group_01, group_11])
        _, p_exp2 = stats.ttest_ind(exp2_treatment, exp2_control)

        # Interaction effect
        # If interaction exists: effect of exp1 differs based on exp2
        effect_exp1_when_exp2_control = np.mean(group_10) - np.mean(group_00)
        effect_exp1_when_exp2_treatment = np.mean(group_11) - np.mean(group_01)
        interaction_magnitude = abs(
            effect_exp1_when_exp2_treatment - effect_exp1_when_exp2_control
        )

        return {
            'exp1_significant': p_exp1 < 0.05,
            'exp2_significant': p_exp2 < 0.05,
            'interaction_magnitude': interaction_magnitude,
            'has_interaction': interaction_magnitude > 0.01  # Threshold
        }