ML Training Data Lake

Versioned datasets for reproducible training

Design a petabyte-scale store for machine learning training data where every dataset version is immutable, every model can be traced to the exact data it saw, and GPUs are fed fast enough to stay busy.

Overview

A training data lake is where an ML organization keeps the data its models learn from: web text, images, logs, labels. It looks like any data lake (files in object storage, a catalog on top) but it has two extra jobs. It must make training reproducible: when a model misbehaves six months later, you must be able to say exactly which rows it was trained on. And it must feed GPUs: a large job reads its dataset again and again at tens of gigabytes a second.

The usual failure is a bucket of files that people overwrite in place. The path s3://data/web-text/ meant one thing in March and another in June, so no one can reproduce the March model or prove what it saw. The fix is the same as for code: versions that never change, and a record of which version every training run used.

Questions to ask first

What modalities? — Text and images; tabular labels as small tables beside them
How much data? — About 5 PB now, growing about 2 PB a year
Who uses it? — A few hundred researchers; tens of large training jobs a week
Must old versions stay readable? — Yes, every version a model was trained on, unless a deletion request forces a change
Is there personal data? — Some raw sources may contain it; nothing unscreened may reach training
How fast must a job read? — About 26 GB/s for the largest job (worked out below)

What you will be able to do

Explain how an open table format makes a dataset version immutable and cheap to keep.
Deduplicate billions of documents with MinHash and LSH, and choose the band settings with numbers.
Size the read path so GPUs stay busy, and say why file size matters on object storage.
Record lineage so a bad source can be traced to every model that used it.
Handle a deletion request without pretending versions are still byte-for-byte reproducible.

Key concepts

Raw and curated zones
Open table format (Iceberg)
Snapshots and tags
Immutable versions
MinHash and LSH
PII screening
Lineage
Large sequential shards
Deletion requests

Functional requirements

Land raw data from many sources in a raw zone that training cannot read
Build curated datasets: exact and near-duplicate removal, PII screening, quality filters
Publish a dataset as a numbered version that never changes afterwards
Record, for every training run, the dataset version and code commit it used
Trace any model back to its raw sources, and any source forward to its models
Search a catalog of datasets by modality, language, license, owner and size
Withdraw a version and rebuild it when data must be deleted

Non-functional requirements

A published version is immutable; reading it next year gives the same rows (unless withdrawn)
Sustain about 26 GB/s of sequential reads for the largest job (a target from the estimate)
New labeled data usable for training within a day (a target we chose)
Unscreened data is never readable by training jobs (enforced by storage permissions)
Catalog search answers in under 1 second (a target we chose)

Interview questions

How do you make dataset versions immutable without copying petabytes?
How do you deduplicate ten billion documents?
How do you keep GPUs from waiting on data?
How do you prove which data a model was trained on?
Where do you screen personal data, and what about deletion requests?
Why not just keep training data in a database?
How do you stop evaluation data leaking into training?