ml-system-design

name	ml-system-design
description	Use when designing end-to-end ML systems, choosing batch vs streaming inference, preventing training/serving skew, building data flywheels, or planning ML infrastructure scaling.

ML System Design

Serving Pattern Decision Table

Factor	Batch	Online (Real-time)	Streaming
Latency tolerance	Hours	<100ms	Seconds
Request pattern	Scheduled/bulk	Per-request	Continuous
Freshness need	Stale OK	Must be fresh	Near-real-time
Compute cost	Low (off-peak)	High (always-on)	Medium
Example	Recommendation email	Search ranking	Fraud detection
Infra complexity	Low	Medium	High
Failure mode	Retry whole job	Per-request retry	Checkpoint + replay

System Design Template

Work through these phases sequentially. Skip none.

Phase 1: Problem and Metrics

1. Business problem → ML problem mapping
   - "Increase engagement" → ranking/recommendation
   - "Reduce fraud" → binary classification
   - "Extract info" → NER/sequence labeling

2. Metrics
   - Business: revenue, CTR, churn rate
   - Model: precision, recall, AUC, NDCG
   - System: p50/p99 latency, throughput, availability

3. Constraints
   - Latency budget (e.g., <50ms for serving)
   - Cost ceiling (e.g., <$0.001 per inference)
   - Data privacy (PII handling, GDPR)

Phase 2: Data Pipeline

Raw Sources → Ingestion → Validation → Feature Engineering → Feature Store
                                |
                          Data Quality Checks
                          - Schema validation
                          - Distribution checks
                          - Freshness checks

# Data validation with Pandera
import pandera as pa

schema = pa.DataFrameSchema({
    "user_id": pa.Column(int, nullable=False),
    "event_ts": pa.Column(pa.DateTime, nullable=False),
    "amount": pa.Column(float, checks=[
        pa.Check.ge(0),
        pa.Check.le(1_000_000),
    ]),
})

validated_df = schema.validate(raw_df)

Phase 3: Feature Engineering

# Shared feature transforms -- used in BOTH training and serving
# This is the single most important thing for preventing skew

class UserFeatures:
    """Feature definitions shared between training and serving."""

    def txn_velocity(self, user_id: str, window_hours: int = 24) -> float:
        """Transaction count in rolling window."""
        # In training: computed from historical log
        # In serving: computed from feature store (pre-aggregated)
        ...

    def avg_order_value(self, user_id: str, window_days: int = 30) -> float:
        ...

    def transform(self, raw: dict) -> dict:
        """Single transform function used everywhere."""
        return {
            "txn_velocity_24h": self.txn_velocity(raw["user_id"], 24),
            "avg_order_30d": self.avg_order_value(raw["user_id"], 30),
            "hour_of_day": raw["timestamp"].hour,
            "day_of_week": raw["timestamp"].weekday(),
            "log_amount": math.log1p(raw["amount"]),
        }

Phase 4: Model Selection

Problem Type	Start With	Scale To
Tabular classification	XGBoost	XGBoost (usually stays best)
Tabular regression	XGBoost	XGBoost + feature engineering
Text classification	Fine-tuned BERT-small	Distilled model from larger LLM
Ranking	LambdaMART	Two-tower + cross-encoder reranker
Recommendation	Matrix factorization	Two-tower retrieval + ranking
Time series	Prophet / statsmodels	Temporal fusion transformer

Phase 5: Training Pipeline

# Training pipeline skeleton
from dataclasses import dataclass

@dataclass
class TrainConfig:
    experiment_name: str
    data_version: str          # DVC or artifact version
    model_type: str
    hyperparams: dict
    train_split_date: str      # Time-based split, never random for time data
    val_split_date: str
    test_split_date: str

def train(config: TrainConfig):
    # 1. Load versioned data
    df = load_data(config.data_version)

    # 2. Time-based split (NEVER random for temporal data)
    train = df[df.date < config.train_split_date]
    val = df[(df.date >= config.val_split_date) & (df.date < config.test_split_date)]
    test = df[df.date >= config.test_split_date]

    # 3. Transform with shared feature code
    features = UserFeatures()
    X_train = train.apply(features.transform, axis=1)

    # 4. Train
    model = fit_model(config.model_type, X_train, train.label, config.hyperparams)

    # 5. Evaluate on test
    metrics = evaluate(model, X_test, test.label)

    # 6. Log everything
    log_experiment(config, metrics, model)

    return model, metrics

Training/Serving Skew Prevention

The #1 cause of ML system failures in production.

Skew Type	Cause	Prevention
Feature skew	Different code paths	Shared transform library
Data skew	Different data sources	Feature store (Feast, Tecton)
Label skew	Label definition change	Versioned label pipelines
Time skew	Future data leakage	Strict time-based splits

# Feature store pattern (Feast example)
from feast import FeatureStore

store = FeatureStore(repo_path="feature_repo/")

# Training: get historical features (point-in-time correct)
training_df = store.get_historical_features(
    entity_df=entity_df_with_timestamps,
    features=["user_features:txn_velocity_24h", "user_features:avg_order_30d"],
).to_df()

# Serving: get latest features (same definitions, same code)
online_features = store.get_online_features(
    features=["user_features:txn_velocity_24h", "user_features:avg_order_30d"],
    entity_rows=[{"user_id": "abc123"}],
).to_dict()

Data Flywheel Design

User Interaction
    → Collect implicit labels (clicks, conversions, dwell time)
    → Log predictions + actuals
    → Identify hard examples (low-confidence predictions)
    → Human review queue (prioritized by uncertainty)
    → Curated labels added to training set
    → Retrain on expanded dataset
    → Deploy improved model
    → Better predictions → more user engagement
    → Repeat

# Hard example mining for the flywheel
def find_hard_examples(predictions: list[dict], threshold: float = 0.1) -> list[dict]:
    """Find predictions where model was least confident."""
    hard = []
    for pred in predictions:
        confidence = abs(pred["score"] - 0.5)  # Distance from decision boundary
        if confidence < threshold:
            hard.append({
                "input": pred["input"],
                "predicted": pred["label"],
                "confidence": pred["score"],
                "needs_review": True,
            })
    return sorted(hard, key=lambda x: x["confidence"])

Scaling Patterns

Horizontal Serving

# Model serving with batched inference (key optimization)
import asyncio
from collections import deque

class BatchingPredictor:
    """Collect individual requests, batch them for GPU efficiency."""

    def __init__(self, model, max_batch: int = 32, max_wait_ms: float = 10):
        self.model = model
        self.max_batch = max_batch
        self.max_wait = max_wait_ms / 1000
        self._queue: deque = deque()
        self._loop = asyncio.get_event_loop()
        self._loop.create_task(self._process_loop())

    async def predict(self, input_data: dict) -> dict:
        future = self._loop.create_future()
        self._queue.append((input_data, future))
        return await future

    async def _process_loop(self):
        while True:
            await asyncio.sleep(self.max_wait)
            if not self._queue:
                continue
            batch = []
            futures = []
             ._queue  (batch) < .max_batch:
                data, fut = ._queue.popleft()
                batch.append(data)
                futures.append(fut)
            
            results = .model.predict_batch(batch)
             fut, result  (futures, results):
                fut.set_result(result)

Model Parallelism Decision

Model Size	Strategy
<2GB	Single GPU, replicate horizontally
2-10GB	Single GPU (A100 80GB), replicate
10-80GB	Tensor parallelism across GPUs
>80GB	Pipeline + tensor parallelism

Cost Modeling

Training cost estimate:
  GPU hours = (dataset_size / batch_throughput) * epochs * hyperparam_trials
  Cost = GPU_hours * hourly_rate

  Example: 10M rows, 50k rows/sec on A100, 10 epochs, 20 trials
  = (10M/50k) * 10 * 20 = 40,000 seconds = ~11 GPU hours
  = 11 * $3.50/hr (A100 spot) = ~$39 per full sweep

Inference cost estimate:
  Cost per request = (GPU_cost_per_second * latency_seconds) / batch_size

  Example: $3.50/hr A100, 20ms latency, batch size 16
  = ($0.00097/sec * 0.02) / 16 = $0.0000012 per request
  = ~$1.20 per million requests

Monitoring

# Model monitoring checks (run on every prediction batch)
from dataclasses import dataclass

@dataclass
class MonitoringAlert:
    metric: str
    current: float
    threshold: float
    severity: str  # "warning" | "critical"

def monitor_predictions(predictions: list[dict], reference_stats: dict) -> list[MonitoringAlert]:
    alerts = []

    # 1. Feature drift (PSI or KS test)
    for feature_name, ref_dist in reference_stats["features"].items():
        current_dist = compute_distribution(predictions, feature_name)
        psi = compute_psi(ref_dist, current_dist)
        if psi > 0.2:
            alerts.append(MonitoringAlert(f"feature_drift:{feature_name}", psi, 0.2, "critical"))
        elif psi > 0.1:
            alerts.append(MonitoringAlert(f"feature_drift:{feature_name}", psi, 0.1, "warning"))

    # 2. Prediction drift
    pred_mean = sum(p["score"] for p in predictions) / len(predictions)
    ref_mean = reference_stats["pred_mean"]
    if abs(pred_mean - ref_mean) / ref_mean > 0.15:
        alerts.append(MonitoringAlert(, pred_mean, ref_mean, ))

    
    p99 = compute_percentile([p[]  p  predictions], )
     p99 > reference_stats[]:
        alerts.append(MonitoringAlert(, p99, reference_stats[], ))

     alerts

Gotchas

• Random splits on temporal data: always use time-based splits; random splits leak future info
• Offline metric obsession: a model with 0.01 higher AUC often makes zero business impact; A/B test everything
• Feature store skipped: "we'll add it later" means training/serving skew ships on day 1
• No shadow mode: deploy new models in shadow mode (predict but don't act) before switching traffic
• Retraining without validation: automated retraining must include automated quality gates; never auto-deploy a worse model
• GPU over-provisioning: most tabular models don't need GPUs at serving time; benchmark CPU first
• Ignoring data quality: garbage in, garbage out; 80% of ML work is data; invest in validation early
• Missing feedback loops: if you can't measure real-world outcomes, you can't improve the model
• Serving the training framework: don't serve with PyTorch if you can export to ONNX; inference runtimes are 2-10x faster
• Batch predictions going stale: if batch runs at midnight and user behavior changes by noon, your predictions are wrong