Agent skill

ml-system-design

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.

Stars 163
Forks 31

Install this agent skill to your Project

npx add-skill https://github.com/majiayu000/claude-skill-registry/tree/main/skills/other/other/ml-system-design

SKILL.md

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
python
# 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

python
# 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

python
# 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
python
# 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
python
# 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

python
# 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 = []
            while self._queue and len(batch) < self.max_batch:
                data, fut = self._queue.popleft()
                batch.append(data)
                futures.append(fut)
            # Single batched GPU call
            results = self.model.predict_batch(batch)
            for fut, result in zip(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

python
# 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("prediction_drift", pred_mean, ref_mean, "warning"))

    # 3. Latency
    p99 = compute_percentile([p["latency_ms"] for p in predictions], 99)
    if p99 > reference_stats["p99_threshold_ms"]:
        alerts.append(MonitoringAlert("latency_p99", p99, reference_stats["p99_threshold_ms"], "critical"))

    return 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

Expand your agent's capabilities with these related and highly-rated skills.

Didn't find tool you were looking for?

Be as detailed as possible for better results