We are struggling with reproducibility because our data changes frequently. While the MLflow Model Registry is great for tracking different versions of our weights, it doesn't natively version the raw S3 data. How are you all linking specific dataset versions to a registered model to ensure we can recreate the exact training conditions a year later?
3 answers
While the registry focuses on models, the best practice is to use 'Tags' and 'Artifacts' within the initial run. When you log your model, you should also log a metadata file or a symlink to the specific DVC (Data Version Control) commit or the S3 version ID used. In the MLflow UI, you can then see exactly which data produced that specific model version. Some teams also use the mlflow.data module introduced in recent versions, which allows you to log the dataset object directly, capturing its schema and a profile of the data, which creates a much tighter bond in the lineage graph than just a simple text comment.
Do you think it’s better to store the actual data as an artifact in MLflow, or should we just store a hash of the dataset?
We use a custom tag called data_version_id. It’s simple, searchable in the UI, and lets us query all models trained on a specific slice of data.
That is exactly what we do too! The searchability of tags in the UI is a killer feature for finding models when a specific data batch is later found to be corrupted.
Matthew, storing the actual data as an artifact is usually a bad idea if the dataset is several gigabytes, as it will bloat your MLflow backend. Storing a hash or a pointer to a versioned bucket is much more scalable. This way, you maintain the "truth" of the model's origin without turning your tracking server into a heavy-duty storage warehouse that’s impossible to back up.