System Design Lab

An ML training platform starts as one notebook and grows the lifecycle around it: pipelines, tracking, a cluster, and a registry.

Change training jobs per day, dataset size, cluster GPUs, experiments tracked, model size, and pipeline stages. Toggle data-parallel distributed training and automated retraining. The design moves from a single notebook to a scheduled pipeline, versioned data and feature prep, a distributed training cluster, an experiment tracker, a model registry, and CI/CD that promotes models to serving.

Read the full tutorial

Guided walkthrough

Reason about it one step at a time

Normal evolution scenarios

Click left to right for the intended demo path. Each card changes the workload inputs.

Workload

These are inputs, not preset architecture stages.

Training jobs / day Distinct training runs launched per day across the team (experiments, sweeps, retrains). Dataset size Size of the training dataset that must be versioned and read each run. Cluster GPUs Accelerators available to the training workers. One host tops out around 8 GPUs. Experiments tracked Experiments/runs tracked (cumulative) whose params, metrics, and artifacts must be compared. Model size Trainable parameters. Past ~1B params one GPU can no longer hold model plus optimizer state. Pipeline stages Steps from raw data to a registered model: ingest, validate, features, train, evaluate, register. Dataset versions / month How often the dataset changes; high churn makes lineage and reproducibility hard by hand. Data-parallel distributed training Shard each batch across many GPUs/hosts and all-reduce gradients to train faster. Each replica still holds a full model copy, so this raises throughput but does not fit a model too big for one GPU. Automated retraining + CI/CD Trigger pipelines on new data or schedule, then promote passing models to serving automatically.

Recommended shape

Current architecture path

Authors

ML author writes training code and launches runs from a notebook or the CLI

Orchestration

Pipeline scheduler turns each run into a scheduled, repeatable DAG of stages

CI/CD trigger fires pipelines on new data or schedule and promotes passing models

Data + features

Data versioning pins an exact dataset version so every run is reproducible

Feature prep cleans and transforms raw data into training features

Training cluster

Training worker runs the training loop on the GPUs of one machine

GPU cluster shards batches across many GPUs and all-reduces gradients

Registry + tracking

Experiment tracker logs params, metrics, and artifacts so runs can be compared

Model registry stores versioned models with stage and approval metadata

Serving handoff receives promoted models for online or batch inference

Bottlenecks

Orchestration load

Data + lineage pressure

GPU cluster capacity

Experiment tracking scale

Registry governance

Why this changes

Decision tradeoffs

Pipeline orchestration

Data versioning + lineage

Distributed training

Experiment tracking

Model registry

Automated retraining / CI-CD

Source-backed rules

These are the durable system-design claims behind the model. The exact slider thresholds are deliberately labeled as teaching assumptions.

Verified rule

A model registry manages the lifecycle of versioned models

MLflow tracks params, metrics, and artifacts per run and registers versioned models with stage transitions, which is why tracking and a registry appear once many experiments and models exist.

MLflow Docs

Verified rule

Pipelines make ML workflows portable, scalable, and repeatable

Kubeflow Pipelines orchestrates multi-step ML workflows as DAGs on Kubernetes, the reason ad-hoc scripts become a scheduled pipeline as stages and run volume grow.

Kubeflow Pipelines

Verified rule

Production ML platforms standardize data, training, and deployment components

TFX defines reusable components (ingest, validate, transform, train, evaluate, push) so the lifecycle from data validation to deployment is consistent and automatable.

TFX Guide

Verified rule

Data versioning gives datasets git-like lineage and reproducibility

DVC version-controls large datasets and pipelines so any run can be reproduced from an exact data version, essential once the dataset changes frequently.

DVC Docs

Teaching assumptions

Single-host training tops out around 8 GPUs and roughly 1B parameters before model and optimizer state stop fitting on one accelerator.
Orchestration, tracking, and storage budgets are conservative teaching numbers chosen to make the five scenarios progress, not vendor limits.
Data-parallel distributed training is modeled as a throughput axis (more replicas finish faster); a model larger than one GPU instead needs model sharding (tensor/pipeline parallelism or FSDP/ZeRO), modeled as a separate sharding signal.