Reproducible ML Training Datasets Without Duplicating Data

A reproducibility pattern for ML teams on Apache Iceberg, using the model registry you already have.

Photo by Gianpi Colonna (Hong Kong)

Who this is for: ML teams who need reproducible training sets — it works whether you train on gigabytes or petabytes.

Reproducing the exact dataset a model was trained on is a baseline requirement. Audits demand it. Debugging a regression depends on it. Retraining from a known starting point, or comparing two model versions on equal footing, is impossible without it. ML work has to be able to point at the precise data behind every model: the same rows, the same filters, the same cut of reality the model actually learned from.

What makes this hard is that the data underneath isn’t a static file — it is a query against a live source table, and that table keeps changing. New rows land continuously. Late-arriving records backfill into date ranges you have already trained on. Labels are corrected weeks after the fact. A column is renamed, a type is widened, a schema evolves. Run the same query a month later and it returns a different dataset — not because the query changed, but because the table beneath it did.

The obvious fix, and what it costs

The instinct is to take a permanent snapshot of the dataset by copying it. You run your filters once and write the result to a dedicated table. The training set is now frozen, self-contained, and trivially reproducible. This is full materialisation, and in many situations it is a perfectly good answer. The problem is what it costs at scale. Every baseline you preserve is a full copy of the filtered data — roughly twice the storage of the source rows it draws from, since you now keep both the original and the cut.

Materialisation resolves the reproducibility problem by spending money, and for large or numerous baselines it spends a lot of it.

Cost is the soft version of this constraint. The harder one is that at sufficient scale, copying the data is not expensive — it is impossible. For a meaningful portion of companies who dream about training large models, just copy it is not the worse option — it is not an option at all.

A better solution

Another approach exists, and it comes from a table format itself. The table format in question is Apache Iceberg — an open table format that sits on top of Parquet (or ORC, or Avro) files in object storage and turns a directory of files into a real, transactional table with schema evolution, partition evolution, time travel, and ACID writes. It is now the default lakehouse format at Netflix, Apple, Stripe, Airbnb, Expedia, and many others, and it is supported as a first-class citizen by Spark, Trino, Flink, Snowflake, BigQuery, and most of the modern data stack. The detail of Iceberg that matters for this pattern is the one in the next section: the way it tracks history.

Figure 1 from How Apache Iceberg Actually Works.

I will not describe the inner workings of Apache Iceberg in this article, to know more, start from its official documentation.

The snapshot is already a version

Apache Iceberg tracks every change to a table as an immutable snapshot — a complete, addressable view of the table at a point in time, identified by a snapshot ID. When new writes arrive, they do not alter the data behind an existing snapshot; they produce a new one. The old state remains in place, still queryable, exactly as it was. So the expensive operation under consideration — freezing a copy of the data — has, in a sense, already been performed by the table format. However, a snapshot_idwill usually be set with a retention policy of weeks, not months or years.

Retention takes two things: a durable, named reference to the right snapshot, and a guarantee that it will not be garbage-collected by Iceberg’s ExpireSnapshots mechanism. Iceberg provides both in a single primitive. A tag is a named pointer to a specific snapshot, and a tagged snapshot is protected from ExpireSnapshots. By default, a tag never expires, but a retention policy can be set on the tag itself.

A tag indicates that a snapshot is preserved, but not why — what the cut represents, and what it will be used for. That mapping has to be recorded somewhere, and the natural place is the model registry — MLflow, Weights & Biases, SageMaker Model Registry, internal model cards — whichever it is, it is already storing run metadata, hyperparameters, metrics, and code version per training run. The data binding belongs there too.

There is also a deliberate split of responsibilities here. The Iceberg tag names the cut of data — independent of any model that might use it. Multiple model variants can train on the same cut, and the tag name should not pretend to belong to any one of them. The model registry, in turn, names which models trained on which cut.

Image from the Author: One cut, many models. The tag names what the data is, not who uses it.

The full pattern follows in two pieces. First, a structured tag name on the source table that pins the snapshot and identifies the data cut it represents — a durable, garbage-collection-safe pointer that any number of models can reference. A workable shape is lock__{cut_name}__{timestamp}, e.g. lock__bookings_h1_2024__20240701.

Second, the lock tag is recorded as MLflow run tags on the training run, alongside the metrics and parameters MLflow already tracks. Another useful piece of information that should be tracked as metadata in the model registry should be the SQL predicate that is used to further filter down the data — including the partition column, and potentially other filters applied to the data.

Doing so, we are solving the reproducibility problem with a one-liner that starts from the model.

Image from the Author: Iceberg pins the data, MLflow records what each model did with it, and the lock tag string ties them together.

When the work happens

Three things happen, in order, on opposite sides of the model’s life. Before training, you pin the data. During training, you record the binding. After training — months later, in front of an auditor or a regression — you reproduce. Each step lives in a different place: the Iceberg catalog, the training script, and a one-liner that starts from an MLflow run ID. The rest of this section walks through each one.

Image from the Author: Each step lives in a different system — but the same lock_tag string ties all three together.

Before training: pin the snapshot

The first step happens in the Iceberg catalog, before any training job starts. Capture the current snapshot ID of the source table — this is the specific version of the data the cut will permanently reference. It is used once, to create the tag; after that, the tag name is what everything else refers to.

from pyiceberg.catalog import load_catalog
catalog     = load_catalog('default')
table       = catalog.load_table('training_booking_data')
snapshot_id = table.current_snapshot().snapshot_id   # e.g. 8821045123456789

Then pin the snapshot with a structured lock tag. Creating the tag before any model trains is what makes the whole pattern safe: it closes the window in which a maintenance job could delete the data the training is about to depend on.

ALTER TABLE training_booking_data
CREATE TAG lock__bookings_h1_2024__20240701
AS OF VERSION 8821045123456789 [RETAIN 365 DAYS];

One tag per cut, regardless of how many models will eventually train on it. Train and validation splits read the same snapshot through different filters, and multiple model variants can read the same snapshot too; the snapshot is the data-side identity, and the per-consumer distinctions belong in the WHERE clause and on the MLflow run. Iceberg tags carry no key-value metadata of their own — they are pointers to a snapshot ID and nothing more. The executable filters live on the MLflow run.

During training: record the binding in MLflow

The second step happens inside the training script. The job reads the tagged snapshot using partition filters and any other SQL predicate necessary. These are then recorded as MLflow run tags. These are MLflow run tags, not Iceberg tags — same word, different namespace.

import mlflow
from pyiceberg.catalog import load_catalog

LOCK_TAG = 'lock__bookings_h1_2024__20240701'
TRAIN_FILTER = """
booking_date BETWEEN '2024-01-01' AND '2024-06-30'
  AND label IS NOT NULL
  AND traffic_scope = 'web'
"""
VAL_FILTER = """
booking_date BETWEEN '2024-07-01' AND '2024-07-31'
  AND label IS NOT NULL
  AND traffic_scope = 'web'
"""
train_df = (spark.read
            .option('tag', LOCK_TAG)
            .table('training_booking_data')
            .filter(TRAIN_FILTER))
val_df   = (spark.read
            .option('tag', LOCK_TAG)
            .table('training_booking_data')
            .filter(VAL_FILTER))
# ... train the model ...
with mlflow.start_run():
    mlflow.set_tag('data.source_table', 'training_booking_data')
    mlflow.set_tag('data.lock_tag',     LOCK_TAG)
    mlflow.set_tag('data.train_filter', TRAIN_FILTER.strip())
    mlflow.set_tag('data.val_filter',   VAL_FILTER.strip())
    mlflow.log_model(...)

The MLflow run now carries everything needed to reconstruct its training data. The Iceberg tag identifies the snapshot. The filter SQL identifies each cut.

So far we have described a simple case of reading only one table; the same logic applies if we’re reading, along with a fact table and other dimension/feature tables that need to be joined with each other. Each one of the tables will be independently tagged and we will record each of the version in MLFlow registry. Each step lives in a different system — but the same lock_tag string ties all three together

After training: reproduce from the run

The third step happens whenever someone needs the original data back — for an audit, a debug session, a head-to-head comparison. Reproduction starts from a run ID. You ask MLflow what data the run saw, and the run tags tell you both halves.

import mlflow

def reproduce_training_data(run_id: str, split: str = 'train'):
    run = mlflow.get_run(run_id)
    tag    = run.data.tags['data.lock_tag']
    filter = run.data.tags[f'data.{split}_filter']
    table  = run.data.tags['data.source_table']
    return (spark.read
            .option('tag', tag)
            .table(table)
            .filter(filter))

train_df = reproduce_training_data(run_id, 'train')
val_df   = reproduce_training_data(run_id, 'val')

The dataset that comes back is identical to the one defined months earlier, because the tag resolves to a snapshot that has not moved.

Querying the pattern

The same metadata supports the questions you ask between training and reproduction.

To go from a tag back to the models trained on it — “which runs used lock__bookings_h1_2024__20240701?" — mlflow.search_runs answers directly. This is also how the question "are any models still using this cut?" gets answered before dropping a tag.

mlflow.search_runs(
    filter_string="tags.`data.lock_tag` = 'lock__bookings_h1_2024__20240701'"
)

To answer “which cuts exist on this table?” — for inventory, audit, or cleanup — list the lock tags on the source. The tag-name convention makes the list directly meaningful:

SELECT name FROM training_booking_data.refs
WHERE  type = 'TAG'
  AND  name LIKE 'lock__%';

A tiny parsing helper turns the names back into structured fields when you need to filter or group:

def parse_lock_tag(name: str) -> dict | None:
    if not name.startswith('lock__'):
        return None
    _, cut_name, timestamp = name.split('__')
    return {'cut_name': cut_name, 'timestamp': timestamp}

The Iceberg tag is the join key between the two sides — the data cut and the models that trained on it — and both sides query naturally without a third system to mediate them.

The two failure modes

• Tag retention misconfiguration. If a retention is set and any model trained on the cut is still in service when that window elapses, the tag expires, the snapshot becomes eligible for deletion, and the next maintenance run can remove it. An intelligent way to manage cleanup is to make it follow from model side, not the data side. A tag is safe to drop when no model still depends on the cut — which mlflow.search_runs answers directly by listing every run whose data.lock_tag matches. Once that list is empty (or contains only retired runs), the data has no remaining consumers and the tag can be released intentionally:

ALTER TABLE training_booking_data DROP TAG lock__bookings_h1_2024__20240701;

• The second failure mode is loss of MLflow run metadata. Treat the MLflow tracking server with the same care you treat the model weights it points at: back up the backend store, version-control the deployment, and do not use a single MLflow instance as both production registry and scratch experiment tracker.

Conclusion

Image from the Author: Two patterns, the same outcome, different tradeoffs.

Materialisation is the right call when the cost of a second copy is worth the unconditional isolation. The Iceberg-tag-plus-MLflow pattern is the right call when storage is the binding constraint and you already use a model registry — which is to say, almost any serious ML team. The pattern adds nothing new to your stack; it makes the systems you already run carry one more piece of information.

If you do not use MLflow specifically, the same shape works with any model registry that supports per-run key-value metadata: Weights & Biases tags, SageMaker Model Registry custom metadata, Vertex AI metadata, internal model cards in Git.

Image from the Author

FAQ

What if my cut needs more than a SQL filter — heavy preprocessing, tokenisation, image decoding, feature transforms?

The Iceberg tag still holds the data version; what changes is the layer between the tagged read and the GPU. As long as the cut itself can be expressed as a SQL predicate against a single source table — no joins, no window functions, no cross-table logic — that last-mile preprocessing can run on CPU in batches with Ray Data, which streams transformed batches to GPUs without ever materialising the full processed dataset. Pinterest documents using Ray Data for exactly this purpose (Last Mile Data Processing with Ray), and the Ray docs list DoorDash, Instacart, Predibase, and ByteDance as further production users.

Careful thought should be given to plan and map where each stage of preprocessing happens. How the data architecture would look like and which libraries should be used. Designing Data-Intensive Applications is probably the most authoritative place to start learning about Data Engineering.

How does this work at petabyte scale?

Hyperscalers (Meta, Netflix, ByteDance, Uber, and Expedia) all converge on the same pattern: lakehouse table formats as source of truth on object storage, snapshot/tag-style pointers for reproducibility, and a CPU preprocessing tier that streams to GPU rather than materialising.

If you want the full picture with citations, these are the primary engineering disclosures:

• Meta — *Data ingestion and Dataloader for Hyper-Scale Recommendation Models* (2022): Tectonic exabyte-scale distributed FS, Hive/DWRF format, disaggregated Data PreProcessing tier with tens to hundreds of CPU nodes per trainer.
• Netflix — *Supporting Diverse ML Systems at Netflix (2024) and [Incremental Processing using Netflix Maestro and Apache Iceberg](https://netflixtechblog.com/incremental-processing-using-netflix-maestro-and-apache-iceberg-b8ba072ddeeb)* (2023): S3 + Iceberg, Metaflow Fast Data + Apache Arrow, ICDC pattern that creates lightweight snapshots referencing existing files without copying.
• ByteDance — *Building EB-level Data Lake using Hudi at ByteDance* (2021): single Hudi tables over 400 PB, 1,000–10,000+ columns, served via column pruning and predicate pushdown.
• Uber — *Hudi incremental reads* (2025): incremental read scoped to checkpoints, replacing 180-day rescans.
• Streaming dataloaders — Ray Data, MosaicML StreamingDataset (deterministic sample ordering independent of device count), NVIDIA DALI (GPU-accelerated last-mile preprocessing), Apache Iceberg branching and tagging.

What about OpenAI, Anthropic, and other foundation-model labs?

There is no primary engineering disclosure from these organisations about how they store, version, or stream pretraining data at scale. Inferring their practice from FAANG and recommender practice is plausible but cannot be directly sourced.