ML Training Pipeline Diagram

An ML training pipeline diagram shows every stage a model passes through on its way to production and back again. This sample is a customer-churn model retrained continuously: data is validated, turned into features, trained, checked against the current champion, registered, deployed and monitored. Drift in production closes the loop by triggering a new run. It is a pipeline (flow) diagram rather than UML; the variant redraws it as a UML activity diagram.

1. The Diagram

Churn Model: Training Pipeline (Continuous Training) Raw data warehouse snapshot Validate data schema + ranges Features joins, encoding Train XGBoost, 5-fold CV Evaluate AUC vs. champion Register model v42 + lineage Deploy canary 5% then 100% Monitor drift, latency, AUC pass fail: tune drift alert / weekly schedule triggers a retrain Feature store offline + online Experiment tracker params, metrics Model registry versions, stage shared stores: Features write the feature store, Train logs runs, Register writes the registry

The top row prepares and trains; the bottom row runs right to left from the gate back to monitoring, so the retraining loop is a short vertical arrow.

  1. Raw data: a dated snapshot of the warehouse tables, so every run is reproducible.
  2. Validate data: schema, null-rate and range checks (for example with Great Expectations); a bad batch stops here instead of producing a bad model.
  3. Features: joins and encodings, written to the feature store so training and serving compute features the same way.
  4. Train: gradient-boosted trees with 5-fold cross-validation; parameters and metrics go to the experiment tracker.
  5. Evaluate is a gate: the challenger must beat the current champion's AUC on the same hold-out set. On failure, the red loop sends it back for tuning.
  6. Register stores the model with its lineage: data snapshot, code commit and metrics.
  7. Deploy serves it to 5% of traffic first, then 100% if error rates hold.
  8. Monitor tracks input drift (PSI), latency and, once labels arrive, real AUC. A drift alert or the weekly schedule triggers the next run.

2. How to Read It

Notation Used Stage Gate Monitor Store flow failure loop trigger / metadata A gate blocks promotion unless its metric check passes.

Blue boxes are automated stages, the yellow box is a quality gate, and green marks the stage that watches the live model. Cylinders are the shared MLOps stores that make runs reproducible.

3. Variant: UML Activity Diagram

Variant: The Same Pipeline as a UML Activity Diagram Data team ML engineer Platform Ingest Validate schema Train + tune [AUC beats champion] Register [else] Deploy canary Notify team

The UML activity version adds swimlanes for ownership and uses formal control nodes: a filled initial node, rounded actions, a decision diamond with guards in square brackets, and a bullseye final node. It is better for agreeing who does what; the pipeline diagram above is better for showing stores and the retraining loop.

4. Pros, Cons, and When to Use

Pros

Cons

When to use

When not to use

5. As Diagrams-as-Code

The main diagram as a Mermaid flowchart:

flowchart LR
  Raw[Raw data snapshot] --> Val[Validate data] --> Feat[Features] --> Train[Train XGBoost CV]
  Train --> Eval{Evaluate vs champion}
  Eval -- pass --> Reg[Register model] --> Dep[Deploy canary 5% then 100%] --> Mon[Monitor drift, latency, AUC]
  Eval -- fail --> Train
  Mon -. drift alert / weekly .-> Raw
  Feat -.-> FS[(Feature store)]
  Train -.-> ET[(Experiment tracker)]
  Reg -.-> MR[(Model registry)]