# System Design: ML Training and Serving Pipeline — Feature Store, Training, and Inference at Scale (2025)

Source: https://www.techinterview.org/post/3233466263/system-design-ml-pipeline/
Updated: 2026-04-17 · techinterview.org

## ML Platform Architecture Overview

An ML platform has three distinct phases: **Feature Engineering** — transform raw data into model-ready features and store them for reuse. **Training Pipeline** — fetch features, train models, evaluate, version, and register. **Serving Pipeline** — retrieve features in real time, run inference, log predictions for feedback loops. The critical insight: the training-serving skew problem. If features are computed differently at training time vs. serving time, the model sees different distributions and performance degrades. A Feature Store with shared computation logic solves this.

## Feature Store: Online and Offline

The Feature Store has two planes: **Offline Store**: a data warehouse (Hive, BigQuery, Snowflake) containing historical feature values with point-in-time correctness (no data leakage — features only use data available at the timestamp of the training label). Used for training data generation. **Online Store**: a low-latency key-value store (Redis, DynamoDB, Cassandra) serving features at inference time with < 10ms p99. Feature entities (user, item, session) are the primary keys.


```
class FeatureStore:
    def get_training_data(self, entity_ids: list, feature_names: list,
                          start_time: datetime, end_time: datetime) -> pd.DataFrame:
        # Point-in-time join: for each label timestamp,
        # fetch the feature value that was current at that time
        query = """
            SELECT label.entity_id, label.timestamp, label.label,
                   feat.value
            FROM labels label
            JOIN feature_log feat
              ON label.entity_id = feat.entity_id
             AND feat.timestamp = (
                 SELECT MAX(timestamp) FROM feature_log
                 WHERE entity_id = label.entity_id
                   AND feature_name = %s
                   AND timestamp  dict:
        keys = [f"{name}:{entity_id}" for name in feature_names]
        values = self.redis.mget(keys)
        return dict(zip(feature_names, values))
```


## Training Pipeline

Orchestrated by a workflow engine (Airflow, Kubeflow, Metaflow). Stages: **Data Validation** — check for schema drift, null rates, feature distribution shifts (Great Expectations or TFX). **Feature Generation** — point-in-time join of labels with offline features. **Model Training** — distributed training on GPU cluster (Kubernetes + GPU nodes); hyperparameter tuning via Ray Tune or Optuna. **Evaluation** — compute held-out metrics; compare against production model (challenger/champion). **Model Registration** — if challenger beats champion by threshold, register in Model Registry (MLflow) with metadata: metrics, training data hash, feature schema version, code SHA. **Deployment trigger** — CI/CD pipeline picks up the new registered model version.

## Model Serving: Online Inference


```
class InferenceService:
    def predict(self, request: PredictRequest) -> PredictResponse:
        # 1. Fetch real-time features from online store
        features = self.feature_store.get_online_features(
            entity_id=request.entity_id,
            feature_names=self.model.required_features
        )

        # 2. Optionally join with request-time context features
        features.update(request.context_features)

        # 3. Run model inference
        score = self.model.predict(features)

        # 4. Log prediction for training feedback loop
        self.prediction_logger.log(PredictionLog(
            entity_id=request.entity_id,
            features=features,
            score=score,
            model_version=self.model.version,
            timestamp=datetime.utcnow()
        ))
        return PredictResponse(score=score)
```

Latency budget for online inference: feature retrieval < 5ms, model forward pass < 10ms, total < 20ms p99. Achieve with: model quantization (INT8), ONNX export + TensorRT, batch inference for non-real-time use cases. **Shadow mode**: run new model version in parallel without serving its results to users — compare predictions to production model to catch regressions before cutover.

## Monitoring: Training-Serving Skew and Data Drift

Log feature distributions at serving time. Compare to training data distributions using PSI (Population Stability Index): PSI = sum((actual% - expected%) * ln(actual% / expected%)). PSI  0.25: significant drift, retrain. **Model performance monitoring**: track prediction distribution (score distribution shift often precedes metric degradation), and compare against ground truth labels once available (delayed labels). Set up automated retraining triggers when drift metrics exceed thresholds. **Feature freshness**: monitor Redis TTL and feature update lag — stale features silently degrade model quality without raising errors.

See also: [Databricks Interview Prep](https://techinterview.org/databricks/)

See also: [Meta Interview Prep](https://techinterview.org/meta/)

See also: [Netflix Interview Prep](https://techinterview.org/netflix/)