Skip to content

Latest commit

 

History

History
518 lines (429 loc) · 15.9 KB

File metadata and controls

518 lines (429 loc) · 15.9 KB

Architecture & Design Document

System Overview

The ML Observability System is a production-grade platform for monitoring machine learning models in production. It provides real-time drift detection, performance monitoring, and alerting capabilities.

High-Level Architecture

┌─────────────────────────────────────────────────────────────────┐
│                     Production ML Models                         │
│                   (Various frameworks/languages)                 │
└────────────────────────┬────────────────────────────────────────┘
                         │ HTTP POST /predictions
                         ▼
┌─────────────────────────────────────────────────────────────────┐
│                        Load Balancer                             │
│                      (Nginx/ALB/etc.)                            │
└────────────────────────┬────────────────────────────────────────┘
                         │
                         ▼
┌─────────────────────────────────────────────────────────────────┐
│                   FastAPI Application Layer                      │
│  ┌───────────────────────────────────────────────────────────┐  │
│  │               API Gateway (FastAPI)                       │  │
│  │  - Request validation (Pydantic)                          │  │
│  │  - Authentication/Authorization                           │  │
│  │  - Rate limiting                                          │  │
│  │  - CORS handling                                          │  │
│  └─────────┬────────────────────────────────────────────────┘  │
│            │                                                     │
│  ┌─────────▼────────────────────────────────────────────────┐  │
│  │              Business Logic Layer                        │  │
│  ├──────────────────────────────────────────────────────────┤  │
│  │  Ingestion Service  │  Drift Detection  │  Performance   │  │
│  │  - Log validation   │  - KS Test        │  - Accuracy    │  │
│  │  - Feature extract  │  - PSI Calc       │  - Precision   │  │
│  │  - Data persistence │  - Chi-Square     │  - Recall      │  │
│  ├──────────────────────────────────────────────────────────┤  │
│  │  Alerting Service   │  Statistics Eng.  │  Dashboard API │  │
│  │  - Threshold check  │  - Aggregations   │  - Metrics     │  │
│  │  - Alert generation │  - Cache mgmt     │  - Trends      │  │
│  │  - Notifications    │  - Query optim.   │  - Summaries   │  │
│  └──────────────────────────────────────────────────────────┘  │
└────────┬───────────────────────────────┬────────────────────────┘
         │                               │
         ▼                               ▼
┌─────────────────────┐        ┌─────────────────────┐
│    PostgreSQL       │        │       Redis         │
│  ┌───────────────┐  │        │  ┌───────────────┐  │
│  │  predictions  │  │        │  │ Drift cache   │  │
│  │  baselines    │  │        │  │ Metrics cache │  │
│  │  drift_results│  │        │  │ Session data  │  │
│  │  perf_metrics │  │        │  │ Rate limits   │  │
│  │  alerts       │  │        │  └───────────────┘  │
│  └───────────────┘  │        └─────────────────────┘
│  - Time-series      │
│  - Partitioning     │
│  - Indexing         │
└─────────────────────┘

         │
         ▼
┌─────────────────────┐
│  Background Jobs    │
│  (Celery Workers)   │
│  - Batch drift det. │
│  - Metric aggreg.   │
│  - Alert delivery   │
│  - Data cleanup     │
└─────────────────────┘

Component Details

1. API Gateway Layer

Technology: FastAPI with Uvicorn

Responsibilities:

  • HTTP request handling
  • Input validation (Pydantic schemas)
  • API documentation (OpenAPI)
  • Error handling and logging
  • Metrics collection (Prometheus)

Endpoints:

POST   /api/v1/predictions           # Log single prediction
POST   /api/v1/predictions/batch     # Batch logging
PATCH  /api/v1/predictions/{id}/gt   # Update ground truth
POST   /api/v1/baselines             # Register baseline
GET    /api/v1/baselines/{model_id}  # Get baseline
GET    /api/v1/drift/{model_id}      # Drift status
GET    /api/v1/performance/{model_id}# Performance metrics
GET    /api/v1/alerts                # List alerts
GET    /health                       # Health check
GET    /metrics                      # Prometheus metrics

2. Business Logic Layer

Ingestion Service

  • Input: Prediction logs from models
  • Processing:
    • Schema validation
    • Feature extraction
    • Data normalization
  • Output: Stored predictions in database
  • Performance: Batch inserts, async processing

Drift Detection Engine

  • Input: Baseline distribution + recent predictions
  • Algorithms:
    • Kolmogorov-Smirnov Test: Numerical features
    • Population Stability Index: Numerical & categorical
    • Chi-Square Test: Categorical features
  • Output: Drift metrics with p-values/PSI scores
  • Optimization: Windowed processing, sampling for high volume

Performance Monitor

  • Input: Predictions with ground truth
  • Metrics:
    • Classification: Accuracy, Precision, Recall, F1
    • Regression: RMSE, MAE, R² (future)
  • Output: Performance trends over time
  • Optimization: Pre-aggregated metrics, time-windowing

Alerting Service

  • Input: Drift results, performance metrics
  • Logic:
    • Threshold evaluation
    • Severity classification
    • Deduplication
  • Output: Alerts to webhooks, email, Slack
  • Features: Acknowledgment, resolution tracking

3. Data Layer

PostgreSQL Database

Schema Design:

-- Predictions (partitioned by timestamp)
predictions (
  id SERIAL PRIMARY KEY,
  model_id VARCHAR(255) NOT NULL,
  timestamp TIMESTAMP NOT NULL,
  features JSONB NOT NULL,
  prediction JSON NOT NULL,
  ground_truth JSON,
  created_at TIMESTAMP DEFAULT NOW()
) PARTITION BY RANGE (timestamp);

-- Indexes
CREATE INDEX idx_predictions_model_ts ON predictions(model_id, timestamp);
CREATE INDEX idx_predictions_gt ON predictions(model_id, ground_truth_timestamp);

-- Baselines
baselines (
  id SERIAL PRIMARY KEY,
  model_id VARCHAR(255) UNIQUE NOT NULL,
  version VARCHAR(100) NOT NULL,
  feature_stats JSONB NOT NULL,
  is_active BOOLEAN DEFAULT true,
  created_at TIMESTAMP,
  updated_at TIMESTAMP
);

-- Drift Results
drift_results (
  id SERIAL PRIMARY KEY,
  model_id VARCHAR(255) NOT NULL,
  feature_name VARCHAR(255) NOT NULL,
  test_type VARCHAR(50) NOT NULL,
  statistic FLOAT NOT NULL,
  p_value FLOAT,
  psi_value FLOAT,
  is_drifted BOOLEAN NOT NULL,
  drift_severity VARCHAR(20) NOT NULL,
  window_start TIMESTAMP NOT NULL,
  window_end TIMESTAMP NOT NULL,
  timestamp TIMESTAMP DEFAULT NOW()
);

CREATE INDEX idx_drift_model_ts ON drift_results(model_id, timestamp);

Partitioning Strategy:

-- Monthly partitions for predictions
CREATE TABLE predictions_2026_01 PARTITION OF predictions
  FOR VALUES FROM ('2026-01-01') TO ('2026-02-01');

CREATE TABLE predictions_2026_02 PARTITION OF predictions
  FOR VALUES FROM ('2026-02-01') TO ('2026-03-01');

Data Retention:

  • Predictions: 90 days (configurable)
  • Drift results: 365 days
  • Aggregated metrics: 730 days
  • Automated cleanup via Celery

Redis Cache

Data Structures:

# Drift status cache
drift_status:{model_id}:{window} -> JSON (TTL: 5 min)

# Performance metrics cache
performance:{model_id}:{window} -> JSON (TTL: 1 hour)

# Baseline cache
baseline:{model_id} -> JSON (TTL: 24 hours)

# Prediction counters (for volume monitoring)
prediction_count:{model_id}:{hour} -> INTEGER (TTL: 7 days)

# Rate limiting
rate_limit:{client_ip}:{endpoint} -> INTEGER (TTL: 1 minute)

4. Background Processing

Celery Tasks:

# Scheduled drift detection (every 15 minutes)
@celery_app.task
def check_drift_all_models():
    for model_id in get_active_models():
        detect_drift(model_id)
        
# Scheduled performance evaluation (hourly)
@celery_app.task
def evaluate_performance():
    for model_id in get_active_models():
        calculate_metrics(model_id)
        
# Data cleanup (daily)
@celery_app.task
def cleanup_old_data():
    delete_old_predictions()
    archive_old_metrics()

Data Flow

Prediction Logging Flow

1. Client → POST /api/v1/predictions
   ↓
2. FastAPI validates request (Pydantic)
   ↓
3. Create Prediction record
   ↓
4. Insert into PostgreSQL
   ↓
5. Increment Redis counter (prediction_count)
   ↓
6. Invalidate drift cache
   ↓
7. Return 201 Created with prediction ID

Drift Detection Flow

1. Client → GET /api/v1/drift/{model_id}
   ↓
2. Check Redis cache
   ├─ Cache hit → Return cached result
   └─ Cache miss ↓
3. Query PostgreSQL for baseline
   ↓
4. Query PostgreSQL for recent predictions (window)
   ↓
5. Extract features from predictions
   ↓
6. For each feature:
   ├─ Numerical → KS Test + PSI
   └─ Categorical → Chi-Square + PSI
   ↓
7. Store drift results in PostgreSQL
   ↓
8. Generate alerts if drift detected
   ↓
9. Cache results in Redis
   ↓
10. Return drift status

Alert Generation Flow

1. Drift/Performance check completes
   ↓
2. Evaluate against thresholds
   ↓
3. If threshold exceeded:
   ├─ Create Alert record
   ├─ Store in PostgreSQL
   ├─ Send to webhook (async)
   └─ Send to email (async)
   ↓
4. Alert stored with:
   - Severity level
   - Detailed context
   - Timestamp

Scalability Design

Horizontal Scaling

API Layer:

Load Balancer
     ├─ API Instance 1
     ├─ API Instance 2
     ├─ API Instance 3
     └─ API Instance N

Stateless Design: No session state in API instances Load Balancing: Round-robin or least-connections Health Checks: /health/ready endpoint

Worker Layer:

Message Queue (Redis)
     ├─ Worker 1 (drift detection)
     ├─ Worker 2 (performance eval)
     ├─ Worker 3 (alerting)
     └─ Worker N (general tasks)

Vertical Scaling Limits

Component CPU RAM Storage Limits
API 4 cores 8 GB - ~10K req/s
PostgreSQL 8 cores 32 GB 1 TB SSD ~50K writes/s
Redis 4 cores 16 GB - ~100K ops/s
Worker 2 cores 4 GB - ~1000 tasks/min

Performance Optimization

Database:

  • Connection pooling (20 connections)
  • Read replicas for dashboard queries
  • Partitioning for time-series data
  • Materialized views for aggregates

Caching:

  • Multi-tier caching (Redis + in-memory)
  • Cache warming for popular models
  • Intelligent cache invalidation

Processing:

  • Batch drift detection (vs. per-request)
  • Sampling for high-volume models
  • Async task processing (Celery)

Security Considerations

Authentication & Authorization

# API Key authentication
@app.middleware("http")
async def verify_api_key(request, call_next):
    api_key = request.headers.get("X-API-Key")
    if settings.require_api_key and not validate_key(api_key):
        return JSONResponse({"error": "Invalid API key"}, status_code=401)
    return await call_next(request)

Input Validation

  • Pydantic schemas for all requests
  • Maximum payload sizes
  • SQL injection prevention (ORM)
  • XSS prevention (sanitization)

Network Security

  • HTTPS/TLS in production
  • CORS configuration
  • Rate limiting per client
  • DDoS protection (via load balancer)

Monitoring & Observability

Metrics (Prometheus)

# Custom metrics
prediction_count = Counter('predictions_total', 'Total predictions')
drift_detected = Counter('drift_detected_total', 'Drift detections')
api_latency = Histogram('api_request_duration_seconds', 'API latency')

Logging

# Structured JSON logging
{
  "timestamp": "2026-01-15T10:30:00Z",
  "level": "INFO",
  "service": "api",
  "model_id": "fraud_detector_v2",
  "action": "drift_detected",
  "details": {
    "feature": "transaction_amount",
    "psi": 0.28,
    "severity": "high"
  }
}

Health Checks

/health        -> Overall system health
/health/live   -> Liveness probe (K8s)
/health/ready  -> Readiness probe (K8s)

Deployment Architecture

Docker Compose (Development)

services:
  api:        # FastAPI application
  postgres:   # PostgreSQL database
  redis:      # Redis cache
  worker:     # Celery worker

Kubernetes (Production)

# Deployments
- api-deployment (3 replicas)
- worker-deployment (2 replicas)

# StatefulSets
- postgres-statefulset (1 primary + 1 replica)
- redis-statefulset (3 nodes, clustered)

# Services
- api-service (LoadBalancer)
- postgres-service (ClusterIP)
- redis-service (ClusterIP)

# Ingress
- nginx-ingress (TLS termination)

Cost Estimation (AWS)

Small Deployment (1K predictions/day):

  • EC2 t3.medium (API): $30/month
  • RDS db.t3.small: $25/month
  • ElastiCache t3.micro: $13/month
  • Total: ~$70/month

Medium Deployment (100K predictions/day):

  • ECS Fargate (API): $150/month
  • RDS db.r5.large: $180/month
  • ElastiCache r5.large: $120/month
  • Total: ~$450/month

Large Deployment (10M predictions/day):

  • EKS cluster: $300/month
  • RDS db.r5.4xlarge: $1,200/month
  • ElastiCache r5.2xlarge: $500/month
  • Load balancer: $20/month
  • Total: ~$2,000/month

Future Enhancements

  1. Advanced Drift Detection

    • Maximum Mean Discrepancy (MMD)
    • Adversarial validation
    • Multivariate tests
  2. Enhanced Monitoring

    • Feature importance drift
    • Model explanation drift (SHAP/LIME)
    • Prediction distribution monitoring
  3. Automation

    • Auto-retraining triggers
    • Self-healing alerts
    • Dynamic threshold adjustment
  4. Integration

    • MLflow integration
    • Kubeflow pipelines
    • SageMaker support
  5. Visualization

    • Real-time dashboard (React/Vue)
    • Interactive charts (Plotly)
    • Custom report generation

Document Version: 1.0.0
Last Updated: January 15, 2026
Maintainer: ML Platform Team