ML System Design
This skill provides frameworks for designing production machine learning systems, from data pipelines to model serving.
When to Use This Skill
Keywords: ML pipeline, machine learning system, feature store, model training, model serving, ML infrastructure, MLOps, A/B testing ML, feature engineering, model deployment
Use this skill when:
- Designing end-to-end ML systems for production
- Planning feature store architecture
- Designing model training pipelines
- Planning model serving infrastructure
- Preparing for ML system design interviews
- Evaluating ML platform tools and frameworks
ML System Architecture Overview
The ML System Lifecycle
┌─────────────────────────────────────────────────────────────────────────┐
│ ML SYSTEM LIFECYCLE │
├─────────────────────────────────────────────────────────────────────────┤
│ │
│ ┌──────────┐ ┌──────────┐ ┌──────────┐ ┌──────────┐ ┌────────┐ │
│ │ Data │──▶│ Feature │──▶│ Model │──▶│ Model │──▶│ Monitor│ │
│ │ Ingestion│ │ Pipeline │ │ Training │ │ Serving │ │ & Eval │ │
│ └──────────┘ └──────────┘ └──────────┘ └──────────┘ └────────┘ │
│ │ │ │ │ │ │
│ ▼ ▼ ▼ ▼ ▼ │
│ ┌──────────┐ ┌──────────┐ ┌──────────┐ ┌──────────┐ ┌────────┐ │
│ │ Data │ │ Feature │ │ Model │ │ Inference│ │ Metrics│ │
│ │ Lake │ │ Store │ │ Registry │ │ Cache │ │ Store │ │
│ └──────────┘ └──────────┘ └──────────┘ └──────────┘ └────────┘ │
│ │
└─────────────────────────────────────────────────────────────────────────┘
Key Components
| Component |
Purpose |
Examples |
| Data Ingestion |
Collect raw data from sources |
Kafka, Kinesis, Pub/Sub |
| Feature Pipeline |
Transform raw data to features |
Spark, Flink, dbt |
| Feature Store |
Store and serve features |
Feast, Tecton, Vertex AI |
| Model Training |
Train and validate models |
SageMaker, Vertex AI, Kubeflow |
| Model Registry |
Version and track models |
MLflow, Weights & Biases |
| Model Serving |
Serve predictions |
TensorFlow Serving, Triton, vLLM |
| Monitoring |
Track model performance |
Evidently, WhyLabs, Arize |
Feature Store Architecture
Why Feature Stores?
Problems without a feature store:
- Training-serving skew (features computed differently)
- Duplicate feature computation across teams
- No feature versioning or lineage
- Slow feature experimentation
Feature Store Components
┌─────────────────────────────────────────────────────────────────┐
│ FEATURE STORE │
├─────────────────────────────────────────────────────────────────┤
│ │
│ ┌─────────────────────┐ ┌─────────────────────┐ │
│ │ OFFLINE STORE │ │ ONLINE STORE │ │
│ │ │ │ │ │
│ │ - Historical data │ │ - Low-latency │ │
│ │ - Training queries │ ────▶ │ - Point lookups │ │
│ │ - Batch features │ sync │ - Real-time serving│ │
│ │ │ │ │ │
│ │ (Data Warehouse) │ │ (Redis, DynamoDB) │ │
│ └─────────────────────┘ └─────────────────────┘ │
│ │
│ ┌─────────────────────────────────────────────────────────────┐│
│ │ FEATURE REGISTRY ││
│ │ - Feature definitions - Version control ││
│ │ - Data lineage - Access control ││
│ └─────────────────────────────────────────────────────────────┘│
└─────────────────────────────────────────────────────────────────┘
Feature Types
| Type |
Computation |
Storage |
Example |
| Batch |
Scheduled (hourly/daily) |
Offline → Online |
User purchase count (30 days) |
| Streaming |
Real-time event processing |
Direct to online |
Items in cart (current) |
| On-demand |
Request-time computation |
Not stored |
Distance to nearest store |
Training-Serving Consistency
TRAINING (Historical):
┌──────────────┐ ┌──────────────┐ ┌──────────────┐
│ Historical │───▶│ Point-in-Time│───▶│ Training │
│ Events │ │ Join │ │ Dataset │
└──────────────┘ └──────────────┘ └──────────────┘
│
Uses feature
definitions
│
SERVING (Real-time): ▼
┌──────────────┐ ┌──────────────┐ ┌──────────────┐
│ Online │───▶│ Same Feature │───▶│ Prediction │
│ Store │ │ Definitions │ │ Request │
└──────────────┘ └──────────────┘ └──────────────┘
Model Training Infrastructure
Training Pipeline Components
┌───────────────────────────────────────────────────────────────────────┐
│ TRAINING PIPELINE │
├───────────────────────────────────────────────────────────────────────┤
│ │
│ ┌────────────┐ ┌────────────┐ ┌────────────┐ ┌────────────┐ │
│ │ Data │──▶│ Feature │──▶│ Model │──▶│ Model │ │
│ │ Loader │ │ Transform│ │ Train │ │ Validate │ │
│ └────────────┘ └────────────┘ └────────────┘ └────────────┘ │
│ │ │ │ │ │
│ ▼ ▼ ▼ ▼ │
│ ┌────────────┐ ┌────────────┐ ┌────────────┐ ┌────────────┐ │
│ │ Experiment │ │ Hyperparameter│ │ Checkpoint │ │ Model │ │
│ │ Tracking │ │ Tuning │ │ Storage │ │ Registry │ │
│ └────────────┘ └────────────┘ └────────────┘ └────────────┘ │
│ │
└───────────────────────────────────────────────────────────────────────┘
Training Infrastructure Patterns
| Pattern |
Use Case |
Tools |
| Single-node |
Small datasets, quick experiments |
Jupyter, local GPU |
| Distributed data-parallel |
Large datasets, same model |
Horovod, PyTorch DDP |
| Model-parallel |
Large models that don't fit in memory |
DeepSpeed, FSDP, Megatron |
| Hyperparameter tuning |
Automated model optimization |
Optuna, Ray Tune |
Experiment Tracking
Track for reproducibility:
| What to Track |
Why |
| Hyperparameters |
Reproduce training runs |
| Metrics |
Compare model performance |
| Artifacts |
Model files, datasets |
| Code version |
Git commit hash |
| Environment |
Docker image, dependencies |
| Data version |
Dataset hash or snapshot |
Model Serving Architecture
Serving Patterns
| Pattern |
Latency |
Throughput |
Use Case |
| Online (REST/gRPC) |
Low (<100ms) |
Medium |
Real-time predictions |
| Batch |
High (hours) |
Very high |
Bulk scoring |
| Streaming |
Medium |
High |
Event-driven predictions |
| Embedded |
Very low |
Varies |
Edge/mobile inference |
Online Serving Architecture
┌─────────────────────────────────────────────────────────────────────┐
│ MODEL SERVING SYSTEM │
├─────────────────────────────────────────────────────────────────────┤
│ │
│ ┌──────────────┐ │
│ │ Clients │ │
│ └──────┬───────┘ │
│ │ │
│ ▼ │
│ ┌──────────────┐ │
│ │ Load Balancer│ │
│ └──────┬───────┘ │
│ │ │
│ ▼ │
│ ┌──────────────────────────────────────────────────────────────┐ │
│ │ API Gateway │ │
│ │ - Authentication - Rate limiting - Request validation │ │
│ └──────────────────────────────┬───────────────────────────────┘ │
│ │ │
│ ┌───────────────────────┼───────────────────────┐ │
│ ▼ ▼ ▼ │
│ ┌────────────┐ ┌────────────┐ ┌────────────┐ │
│ │ Model A │ │ Model B │ │ Model C │ │
│ │ (v1.2) │ │ (v2.0) │ │ (v1.0) │ │
│ └────────────┘ └────────────┘ └────────────┘ │
│ │ │ │ │
│ └───────────────────────┼───────────────────────┘ │
│ ▼ │
│ ┌────────────────┐ │
│ │ Feature Store │ │
│ │ (Online) │ │
│ └────────────────┘ │
│ │
└─────────────────────────────────────────────────────────────────────┘
Latency Optimization
| Technique |
Latency Impact |
Trade-off |
| Batching |
Reduces per-request |
Increases latency for first request |
| Caching |
10-100x faster |
May serve stale predictions |
| Quantization |
2-4x faster |
Slight accuracy loss |
| Distillation |
Variable |
Training overhead |
| GPU inference |
10-100x faster |
Cost increase |
A/B Testing ML Models
Experiment Design
┌─────────────────────────────────────────────────────────────────────┐
│ A/B TESTING ARCHITECTURE │
├─────────────────────────────────────────────────────────────────────┤
│ │
│ ┌──────────────┐ │
│ │ Traffic │ │
│ └──────┬───────┘ │
│ │ │
│ ▼ │
│ ┌──────────────────────┐ │
│ │ Experiment Assignment │ ◀─────── Experiment Config │
│ │ - User bucketing │ - Allocation % │
│ │ - Feature flags │ - Target segments │
│ └──────────┬───────────┘ - Guardrails │
│ │ │
│ ┌────────┴────────┐ │
│ ▼ ▼ │
│ ┌────────┐ ┌────────┐ │
│ │Control │ │Treatment│ │
│ │Model A │ │Model B │ │
│ └────┬───┘ └────┬───┘ │
│ │ │ │
│ └────────┬───────┘ │
│ ▼ │
│ ┌────────────────┐ │
│ │ Metrics Logger │ │
│ └────────┬───────┘ │
│ ▼ │
│ ┌────────────────┐ │
│ │ Statistical │ ─────▶ Decision: Ship / Iterate / Kill │
│ │ Analysis │ │
│ └────────────────┘ │
│ │
└─────────────────────────────────────────────────────────────────────┘
Metrics to Track
| Metric Type |
Examples |
Purpose |
| Model metrics |
AUC, RMSE, precision/recall |
Model quality |
| Business metrics |
CTR, conversion, revenue |
Business impact |
| Guardrail metrics |
Latency, error rate, engagement |
Prevent regressions |
| Segment metrics |
Metrics by user segment |
Detect heterogeneous effects |
Statistical Considerations
- Sample size: Calculate power before experiment
- Duration: Account for novelty effects and time patterns
- Multiple testing: Adjust for multiple metrics (Bonferroni, FDR)
- Early stopping: Use sequential testing methods
Model Monitoring
What to Monitor
| Category |
Metrics |
Alert Threshold |
| Data quality |
Missing values, schema drift |
>1% change |
| Feature drift |
Distribution shift (PSI, KL) |
PSI >0.2 |
| Prediction drift |
Output distribution shift |
Depends on use case |
| Model performance |
Accuracy, AUC (when labels available) |
>5% degradation |
| Operational |
Latency, throughput, errors |
SLO violations |
Drift Detection
┌─────────────────────────────────────────────────────────────────────┐
│ DRIFT DETECTION PIPELINE │
├─────────────────────────────────────────────────────────────────────┤
│ │
│ Training Data Production Data │
│ ┌──────────────┐ ┌──────────────┐ │
│ │ Reference │ │ Current │ │
│ │ Distribution │ │ Distribution │ │
│ └──────┬───────┘ └──────┬───────┘ │
│ │ │ │
│ └──────────────┬──────────────┘ │
│ ▼ │
│ ┌──────────────────┐ │
│ │ Statistical Test │ │
│ │ - PSI (Population Stability Index) │
│ │ - KS Test │
│ │ - Chi-squared │
│ └────────┬─────────┘ │
│ ▼ │
│ ┌──────────────────┐ │
│ │ Drift Score │ │
│ └────────┬─────────┘ │
│ │ │
│ ┌───────────┼───────────┐ │
│ ▼ ▼ ▼ │
│ No Drift Warning Critical │
│ (< 0.1) (0.1-0.2) (> 0.2) │
│ │ │ │ │
│ ▼ ▼ ▼ │
│ Continue Investigate Retrain │
│ │
└─────────────────────────────────────────────────────────────────────┘
Common ML System Design Patterns
Pattern 1: Recommendation System
Components needed:
- Candidate Generation (retrieve 100s-1000s)
- Ranking Model (score and sort)
- Feature Store (user features, item features)
- Real-time personalization (recent behavior)
- A/B testing infrastructure
Pattern 2: Fraud Detection
Components needed:
- Real-time feature computation
- Low-latency model serving (<50ms)
- High recall focus (can't miss fraud)
- Explainability for compliance
- Human-in-the-loop review
- Feedback loop for labels
Pattern 3: Search Ranking
Components needed:
- Two-stage ranking (retrieval + ranking)
- Feature store for query/document features
- Low latency (<200ms end-to-end)
- Learning to rank models
- Click-through rate prediction
- A/B testing with interleaving
Estimation for ML Systems
Training Infrastructure
Training time estimation:
- Dataset size: 100M examples
- Model: Transformer (100M params)
- GPU: A100 (80GB, 312 TFLOPS)
- Batch size: 32
- Training steps: Dataset / batch = 3.1M steps
- Time per step: ~100ms
- Total time: ~86 hours single GPU
- With 8 GPUs (data parallel): ~11 hours
Serving Infrastructure
Inference estimation:
- QPS: 10,000
- Model latency: 20ms
- Batch size: 1 (real-time)
- GPU utilization: 50% (latency constraint)
- Requests per GPU/sec: 25
- GPUs needed: 10,000 / 25 = 400 GPUs
- With batching (batch 8): 100 GPUs (4x reduction)
Related Skills
llm-serving-patterns - LLM-specific serving and optimizationrag-architecture - Retrieval-Augmented Generation patternsvector-databases - Vector search and embeddingsml-inference-optimization - Latency and cost optimizationestimation-techniques - Back-of-envelope calculationsquality-attributes-taxonomy - NFR definitions
Related Commands
/sd:ml-pipeline <problem> - Design ML system interactively/sd:estimate <scenario> - Capacity calculations
Related Agents
ml-systems-designer - Design ML architecturesml-interviewer - Mock ML system design interviews
Version History
- v1.0.0 (2025-12-26): Initial release
Last Updated
Date: 2025-12-26
Model: claude-opus-4-5-20251101